Use tokio_util::CancellationToken everywhere
This commit is contained in:
Igor Katson
parent
53868ad45e
commit
bed7433d8e
16 changed files with 176 additions and 178 deletions
|
|
@ -107,8 +107,7 @@ mod tests {
|
|||
init_logging();
|
||||
|
||||
let info_hash = Id20::from_str("cab507494d02ebb1178b38f2e9d7be299c86b862").unwrap();
|
||||
let (dht, run_dht) = DhtBuilder::new().await.unwrap();
|
||||
tokio::spawn(run_dht);
|
||||
let dht = DhtBuilder::new().await.unwrap();
|
||||
|
||||
let peer_rx = dht.get_peers(info_hash, None).unwrap();
|
||||
let peer_id = generate_peer_id();
|
||||
|
|
|
|||
|
|
@ -21,6 +21,7 @@ use librqbit_core::{
|
|||
directories::get_configuration_directory,
|
||||
magnet::Magnet,
|
||||
peer_id::generate_peer_id,
|
||||
spawn_utils::spawn_with_cancel,
|
||||
torrent_metainfo::{torrent_from_bytes, TorrentMetaV1Info, TorrentMetaV1Owned},
|
||||
};
|
||||
use parking_lot::RwLock;
|
||||
|
|
@ -32,12 +33,13 @@ use tokio::{
|
|||
io::AsyncReadExt,
|
||||
net::{TcpListener, TcpStream},
|
||||
};
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tracing::{debug, error, error_span, info, trace, warn, Instrument};
|
||||
|
||||
use crate::{
|
||||
dht_utils::{read_metainfo_from_peer_receiver, ReadMetainfoResult},
|
||||
peer_connection::{with_timeout, PeerConnectionOptions},
|
||||
spawn_utils::{spawn, BlockingSpawner},
|
||||
spawn_utils::BlockingSpawner,
|
||||
torrent_state::{
|
||||
ManagedTorrentBuilder, ManagedTorrentHandle, ManagedTorrentState, TorrentStateLive,
|
||||
},
|
||||
|
|
@ -150,23 +152,6 @@ struct SerializedSessionDatabase {
|
|||
torrents: HashMap<usize, SerializedTorrent>,
|
||||
}
|
||||
|
||||
fn spawn_with_cancel_token(
|
||||
mut cancel_rx: tokio::sync::watch::Receiver<()>,
|
||||
name: &str,
|
||||
span: tracing::Span,
|
||||
fut: impl std::future::Future<Output = anyhow::Result<()>> + Send + 'static,
|
||||
) {
|
||||
spawn(name, span, async move {
|
||||
tokio::select! {
|
||||
r = fut => r,
|
||||
_ = cancel_rx.changed() => {
|
||||
debug!("task canceled");
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
pub struct Session {
|
||||
peer_id: Id20,
|
||||
dht: Option<Dht>,
|
||||
|
|
@ -178,8 +163,7 @@ pub struct Session {
|
|||
|
||||
tcp_listen_port: Option<u16>,
|
||||
|
||||
cancel_tx: tokio::sync::watch::Sender<()>,
|
||||
cancel_rx: tokio::sync::watch::Receiver<()>,
|
||||
cancellation_token: CancellationToken,
|
||||
}
|
||||
|
||||
async fn torrent_from_url(url: &str) -> anyhow::Result<TorrentMetaV1Owned> {
|
||||
|
|
@ -395,14 +379,17 @@ impl Session {
|
|||
Ok(dir.data_dir().join("session.json"))
|
||||
}
|
||||
|
||||
pub fn cancellation_token(&self) -> &CancellationToken {
|
||||
&self.cancellation_token
|
||||
}
|
||||
|
||||
/// Create a new session with options.
|
||||
pub async fn new_with_opts(
|
||||
output_folder: PathBuf,
|
||||
mut opts: SessionOptions,
|
||||
) -> anyhow::Result<Arc<Self>> {
|
||||
let peer_id = opts.peer_id.unwrap_or_else(generate_peer_id);
|
||||
|
||||
let (cancel_tx, cancel_rx) = tokio::sync::watch::channel(());
|
||||
let token = CancellationToken::new();
|
||||
|
||||
let (tcp_listener, tcp_listen_port) = if let Some(port_range) = opts.listen_port_range {
|
||||
let (l, p) = create_tcp_listener(port_range)
|
||||
|
|
@ -418,24 +405,17 @@ impl Session {
|
|||
None
|
||||
} else {
|
||||
let dht = if opts.disable_dht_persistence {
|
||||
let (dht, run_worker) = DhtBuilder::with_config(DhtConfig::default())
|
||||
.await
|
||||
.context("error initializing DHT")?;
|
||||
spawn_with_cancel_token(cancel_rx.clone(), "dht", error_span!("dht"), run_worker);
|
||||
dht
|
||||
DhtBuilder::with_config(DhtConfig {
|
||||
cancellation_token: Some(token.child_token()),
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.context("error initializing DHT")?
|
||||
} else {
|
||||
let pdht_config = opts.dht_config.take().unwrap_or_default();
|
||||
let (dht, run_worker, run_persistence) = PersistentDht::create(Some(pdht_config))
|
||||
PersistentDht::create(Some(pdht_config), Some(token.clone()))
|
||||
.await
|
||||
.context("error initializing persistent DHT")?;
|
||||
spawn_with_cancel_token(cancel_rx.clone(), "dht", error_span!("dht"), run_worker);
|
||||
spawn_with_cancel_token(
|
||||
cancel_rx.clone(),
|
||||
"dht_persistence",
|
||||
error_span!("dht_persistence"),
|
||||
run_persistence,
|
||||
);
|
||||
dht
|
||||
.context("error initializing persistent DHT")?
|
||||
};
|
||||
|
||||
Some(dht)
|
||||
|
|
@ -455,14 +435,12 @@ impl Session {
|
|||
spawner,
|
||||
output_folder,
|
||||
db: RwLock::new(Default::default()),
|
||||
cancel_rx,
|
||||
cancel_tx,
|
||||
cancellation_token: token,
|
||||
tcp_listen_port,
|
||||
});
|
||||
|
||||
if let Some(tcp_listener) = tcp_listener {
|
||||
session.spawn(
|
||||
"tcp listener",
|
||||
error_span!("tcp_listen", port = tcp_listen_port),
|
||||
session.clone().task_tcp_listener(tcp_listener),
|
||||
);
|
||||
|
|
@ -471,7 +449,6 @@ impl Session {
|
|||
if let Some(listen_port) = tcp_listen_port {
|
||||
if opts.enable_upnp_port_forwarding {
|
||||
session.spawn(
|
||||
"upnp_forward",
|
||||
error_span!("upnp_forward", port = listen_port),
|
||||
session.clone().task_upnp_port_forwarder(listen_port),
|
||||
);
|
||||
|
|
@ -489,11 +466,7 @@ impl Session {
|
|||
})?;
|
||||
}
|
||||
let persistence_task = session.clone().task_persistence();
|
||||
session.spawn(
|
||||
"session persistene",
|
||||
error_span!("session_persistence"),
|
||||
persistence_task,
|
||||
);
|
||||
session.spawn(error_span!("session_persistence"), persistence_task);
|
||||
}
|
||||
|
||||
Ok(session)
|
||||
|
|
@ -645,11 +618,10 @@ impl Session {
|
|||
/// Spawn a task in the context of the session.
|
||||
pub fn spawn(
|
||||
&self,
|
||||
name: &str,
|
||||
span: tracing::Span,
|
||||
fut: impl std::future::Future<Output = anyhow::Result<()>> + Send + 'static,
|
||||
) {
|
||||
spawn_with_cancel_token(self.cancel_rx.clone(), name, span, fut);
|
||||
spawn_with_cancel(span, self.cancellation_token.clone(), fut);
|
||||
}
|
||||
|
||||
/// Stop the session and all managed tasks.
|
||||
|
|
@ -666,7 +638,7 @@ impl Session {
|
|||
debug!("error pausing torrent: {e:#}");
|
||||
}
|
||||
}
|
||||
let _ = self.cancel_tx.send(());
|
||||
self.cancellation_token.cancel();
|
||||
// this sucks, but hopefully will be enough
|
||||
tokio::time::sleep(Duration::from_secs(1)).await;
|
||||
}
|
||||
|
|
@ -999,6 +971,7 @@ impl Session {
|
|||
builder
|
||||
.overwrite(opts.overwrite)
|
||||
.spawner(self.spawner)
|
||||
.cancellation_token(self.cancellation_token.child_token())
|
||||
.peer_id(self.peer_id);
|
||||
|
||||
if opts.disable_trackers {
|
||||
|
|
|
|||
|
|
@ -65,6 +65,7 @@ use itertools::Itertools;
|
|||
use librqbit_core::{
|
||||
id20::Id20,
|
||||
lengths::{ChunkInfo, Lengths, ValidPieceIndex},
|
||||
spawn_utils::spawn_with_cancel,
|
||||
speed_estimator::SpeedEstimator,
|
||||
torrent_metainfo::TorrentMetaV1Info,
|
||||
};
|
||||
|
|
@ -80,6 +81,7 @@ use tokio::{
|
|||
},
|
||||
time::timeout,
|
||||
};
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tracing::{debug, error, error_span, info, trace, warn};
|
||||
use url::Url;
|
||||
|
||||
|
|
@ -90,7 +92,6 @@ use crate::{
|
|||
PeerConnection, PeerConnectionHandler, PeerConnectionOptions, WriterRequest,
|
||||
},
|
||||
session::CheckedIncomingConnection,
|
||||
spawn_utils::spawn,
|
||||
torrent_state::{peer::Peer, utils::atomic_inc},
|
||||
tracker_comms::{TrackerError, TrackerRequest, TrackerRequestEvent, TrackerResponse},
|
||||
type_aliases::{PeerHandle, BF},
|
||||
|
|
@ -185,17 +186,16 @@ pub struct TorrentStateLive {
|
|||
|
||||
finished_notify: Notify,
|
||||
|
||||
cancel_tx: tokio::sync::watch::Sender<()>,
|
||||
cancel_rx: tokio::sync::watch::Receiver<()>,
|
||||
|
||||
down_speed_estimator: SpeedEstimator,
|
||||
up_speed_estimator: SpeedEstimator,
|
||||
cancellation_token: CancellationToken,
|
||||
}
|
||||
|
||||
impl TorrentStateLive {
|
||||
pub(crate) fn new(
|
||||
paused: TorrentStatePaused,
|
||||
fatal_errors_tx: tokio::sync::oneshot::Sender<anyhow::Error>,
|
||||
cancellation_token: CancellationToken,
|
||||
) -> Arc<Self> {
|
||||
let (peer_queue_tx, peer_queue_rx) = unbounded_channel();
|
||||
|
||||
|
|
@ -206,8 +206,6 @@ impl TorrentStateLive {
|
|||
let needed_bytes = paused.info.lengths.total_length() - have_bytes;
|
||||
let lengths = *paused.chunk_tracker.get_lengths();
|
||||
|
||||
let (cancel_tx, cancel_rx) = tokio::sync::watch::channel(());
|
||||
|
||||
let state = Arc::new(TorrentStateLive {
|
||||
meta: paused.info.clone(),
|
||||
peers: Default::default(),
|
||||
|
|
@ -229,20 +227,17 @@ impl TorrentStateLive {
|
|||
finished_notify: Notify::new(),
|
||||
down_speed_estimator,
|
||||
up_speed_estimator,
|
||||
cancel_rx,
|
||||
cancel_tx,
|
||||
cancellation_token,
|
||||
});
|
||||
|
||||
for tracker in state.meta.trackers.iter() {
|
||||
state.spawn(
|
||||
"tracker_monitor",
|
||||
error_span!(parent: state.meta.span.clone(), "tracker_monitor", url = tracker.to_string()),
|
||||
state.clone().task_single_tracker_monitor(tracker.clone()),
|
||||
);
|
||||
}
|
||||
|
||||
state.spawn(
|
||||
"speed_estimator_updater",
|
||||
error_span!(parent: state.meta.span.clone(), "speed_estimator_updater"),
|
||||
{
|
||||
let state = Arc::downgrade(&state);
|
||||
|
|
@ -273,29 +268,18 @@ impl TorrentStateLive {
|
|||
);
|
||||
|
||||
state.spawn(
|
||||
"peer_adder",
|
||||
error_span!(parent: state.meta.span.clone(), "peer_adder"),
|
||||
state.clone().task_peer_adder(peer_queue_rx),
|
||||
);
|
||||
state
|
||||
}
|
||||
|
||||
fn spawn(
|
||||
pub(crate) fn spawn(
|
||||
&self,
|
||||
name: &str,
|
||||
span: tracing::Span,
|
||||
fut: impl std::future::Future<Output = anyhow::Result<()>> + Send + 'static,
|
||||
) {
|
||||
let mut cancel_rx = self.cancel_rx.clone();
|
||||
spawn(name, span, async move {
|
||||
tokio::select! {
|
||||
r = fut => r,
|
||||
_ = cancel_rx.changed() => {
|
||||
debug!("task canceled");
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
});
|
||||
spawn_with_cancel(span, self.cancellation_token.clone(), fut);
|
||||
}
|
||||
|
||||
pub fn down_speed_estimator(&self) -> &SpeedEstimator {
|
||||
|
|
@ -418,7 +402,6 @@ impl TorrentStateLive {
|
|||
atomic_inc(&counters.incoming_connections);
|
||||
|
||||
self.spawn(
|
||||
"incoming peer",
|
||||
error_span!(
|
||||
parent: self.meta.span.clone(),
|
||||
"manage_incoming_peer",
|
||||
|
|
@ -570,7 +553,6 @@ impl TorrentStateLive {
|
|||
|
||||
let permit = state.peer_semaphore.clone().acquire_owned().await?;
|
||||
state.spawn(
|
||||
"manage_peer",
|
||||
error_span!(parent: state.meta.span.clone(), "manage_peer", peer = addr.to_string()),
|
||||
state.clone().task_manage_outgoing_peer(addr, permit),
|
||||
);
|
||||
|
|
@ -682,7 +664,6 @@ impl TorrentStateLive {
|
|||
|
||||
// We don't want to remember this task as there may be too many.
|
||||
self.spawn(
|
||||
"transmit_haves",
|
||||
error_span!(
|
||||
parent: self.meta.span.clone(),
|
||||
"transmit_haves",
|
||||
|
|
@ -744,7 +725,7 @@ impl TorrentStateLive {
|
|||
}
|
||||
|
||||
pub fn pause(&self) -> anyhow::Result<TorrentStatePaused> {
|
||||
let _ = self.cancel_tx.send(());
|
||||
self.cancellation_token.cancel();
|
||||
|
||||
let mut g = self.locked.write();
|
||||
|
||||
|
|
@ -971,7 +952,6 @@ impl PeerHandler {
|
|||
|
||||
if let Some(dur) = backoff {
|
||||
self.state.clone().spawn(
|
||||
"wait_for_peer",
|
||||
error_span!(
|
||||
parent: self.state.meta.span.clone(),
|
||||
"wait_for_peer",
|
||||
|
|
|
|||
|
|
@ -20,12 +20,14 @@ use librqbit_core::id20::Id20;
|
|||
use librqbit_core::lengths::Lengths;
|
||||
use librqbit_core::peer_id::generate_peer_id;
|
||||
|
||||
use librqbit_core::spawn_utils::spawn_with_cancel;
|
||||
use librqbit_core::torrent_metainfo::TorrentMetaV1Info;
|
||||
pub use live::*;
|
||||
use parking_lot::RwLock;
|
||||
|
||||
use tokio::time::timeout;
|
||||
use tokio_stream::StreamExt;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tracing::debug;
|
||||
use tracing::error_span;
|
||||
use tracing::trace;
|
||||
|
|
@ -33,7 +35,6 @@ use tracing::warn;
|
|||
use url::Url;
|
||||
|
||||
use crate::chunk_tracker::ChunkTracker;
|
||||
use crate::spawn_utils::spawn;
|
||||
use crate::spawn_utils::BlockingSpawner;
|
||||
use crate::torrent_state::stats::LiveStats;
|
||||
|
||||
|
|
@ -91,6 +92,7 @@ pub struct ManagedTorrentInfo {
|
|||
|
||||
pub struct ManagedTorrent {
|
||||
pub info: Arc<ManagedTorrentInfo>,
|
||||
pub cancellation_token: CancellationToken,
|
||||
pub(crate) only_files: Option<Vec<usize>>,
|
||||
locked: RwLock<ManagedTorrentLocked>,
|
||||
}
|
||||
|
|
@ -179,10 +181,11 @@ impl ManagedTorrent {
|
|||
let spawn_fatal_errors_receiver =
|
||||
|state: &Arc<Self>, rx: tokio::sync::oneshot::Receiver<anyhow::Error>| {
|
||||
let span = state.info.span.clone();
|
||||
let token = state.cancellation_token.clone();
|
||||
let state = Arc::downgrade(state);
|
||||
spawn(
|
||||
"fatal_errors_receiver",
|
||||
spawn_with_cancel(
|
||||
error_span!(parent: span, "fatal_errors_receiver"),
|
||||
token,
|
||||
async move {
|
||||
let e = match rx.await {
|
||||
Ok(e) => e,
|
||||
|
|
@ -191,7 +194,7 @@ impl ManagedTorrent {
|
|||
if let Some(state) = state.upgrade() {
|
||||
state.stop_with_error(e);
|
||||
} else {
|
||||
warn!("tried to stop the torrent with error, but it's couldn't upgrade the arc");
|
||||
warn!("tried to stop the torrent with error, but couldn't upgrade the arc");
|
||||
}
|
||||
Ok(())
|
||||
},
|
||||
|
|
@ -203,40 +206,42 @@ impl ManagedTorrent {
|
|||
initial_peers: Vec<SocketAddr>,
|
||||
peer_rx: Option<RequestPeersStream>,
|
||||
) {
|
||||
let span = live.meta().span.clone();
|
||||
let live = Arc::downgrade(live);
|
||||
spawn(
|
||||
"external_peer_adder",
|
||||
error_span!(parent: span, "external_peer_adder"),
|
||||
async move {
|
||||
{
|
||||
let live: Arc<TorrentStateLive> =
|
||||
live.upgrade().context("no longer live")?;
|
||||
live.spawn(
|
||||
error_span!(parent: live.meta().span.clone(), "external_peer_adder"),
|
||||
{
|
||||
let live = live.clone();
|
||||
async move {
|
||||
trace!("adding {} initial peers", initial_peers.len());
|
||||
for peer in initial_peers {
|
||||
live.add_peer_if_not_seen(peer).context("torrent closed")?;
|
||||
}
|
||||
}
|
||||
|
||||
let mut peer_rx = if let Some(peer_rx) = peer_rx {
|
||||
peer_rx
|
||||
} else {
|
||||
return Ok(());
|
||||
};
|
||||
let live = {
|
||||
let weak = Arc::downgrade(&live);
|
||||
drop(live);
|
||||
weak
|
||||
};
|
||||
|
||||
loop {
|
||||
match timeout(Duration::from_secs(5), peer_rx.next()).await {
|
||||
Ok(Some(peer)) => {
|
||||
let live = match live.upgrade() {
|
||||
Some(live) => live,
|
||||
None => return Ok(()),
|
||||
};
|
||||
live.add_peer_if_not_seen(peer).context("torrent closed")?;
|
||||
let mut peer_rx = if let Some(peer_rx) = peer_rx {
|
||||
peer_rx
|
||||
} else {
|
||||
return Ok(());
|
||||
};
|
||||
|
||||
loop {
|
||||
match timeout(Duration::from_secs(5), peer_rx.next()).await {
|
||||
Ok(Some(peer)) => {
|
||||
let live = match live.upgrade() {
|
||||
Some(live) => live,
|
||||
None => return Ok(()),
|
||||
};
|
||||
live.add_peer_if_not_seen(peer).context("torrent closed")?;
|
||||
}
|
||||
Ok(None) => return Ok(()),
|
||||
// If timeout, check if the torrent is live.
|
||||
Err(_) if live.strong_count() == 0 => return Ok(()),
|
||||
Err(_) => continue,
|
||||
}
|
||||
Ok(None) => return Ok(()),
|
||||
// If timeout, check if the torrent is live.
|
||||
Err(_) if live.strong_count() == 0 => return Ok(()),
|
||||
Err(_) => continue,
|
||||
}
|
||||
}
|
||||
},
|
||||
|
|
@ -252,9 +257,10 @@ impl ManagedTorrent {
|
|||
drop(g);
|
||||
let t = self.clone();
|
||||
let span = self.info().span.clone();
|
||||
spawn(
|
||||
"initialize_and_start",
|
||||
let token = self.cancellation_token.clone();
|
||||
spawn_with_cancel(
|
||||
error_span!(parent: span.clone(), "initialize_and_start"),
|
||||
token.clone(),
|
||||
async move {
|
||||
match init.check().await {
|
||||
Ok(paused) => {
|
||||
|
|
@ -271,7 +277,7 @@ impl ManagedTorrent {
|
|||
}
|
||||
|
||||
let (tx, rx) = tokio::sync::oneshot::channel();
|
||||
let live = TorrentStateLive::new(paused, tx);
|
||||
let live = TorrentStateLive::new(paused, tx, token.child_token());
|
||||
g.state = ManagedTorrentState::Live(live.clone());
|
||||
|
||||
spawn_fatal_errors_receiver(&t, rx);
|
||||
|
|
@ -292,7 +298,11 @@ impl ManagedTorrent {
|
|||
ManagedTorrentState::Paused(_) => {
|
||||
let paused = g.state.take().assert_paused();
|
||||
let (tx, rx) = tokio::sync::oneshot::channel();
|
||||
let live = TorrentStateLive::new(paused, tx);
|
||||
let live = TorrentStateLive::new(
|
||||
paused,
|
||||
tx,
|
||||
self.cancellation_token.child_token().clone(),
|
||||
);
|
||||
g.state = ManagedTorrentState::Live(live.clone());
|
||||
spawn_fatal_errors_receiver(self, rx);
|
||||
spawn_peer_adder(&live, initial_peers, peer_rx);
|
||||
|
|
@ -409,6 +419,7 @@ pub struct ManagedTorrentBuilder {
|
|||
peer_id: Option<Id20>,
|
||||
overwrite: bool,
|
||||
spawner: Option<BlockingSpawner>,
|
||||
cancellation_token: Option<CancellationToken>,
|
||||
}
|
||||
|
||||
impl ManagedTorrentBuilder {
|
||||
|
|
@ -429,9 +440,15 @@ impl ManagedTorrentBuilder {
|
|||
trackers: Default::default(),
|
||||
peer_id: None,
|
||||
overwrite: false,
|
||||
cancellation_token: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn cancellation_token(&mut self, token: CancellationToken) -> &mut Self {
|
||||
self.cancellation_token = Some(token);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn only_files(&mut self, only_files: Vec<usize>) -> &mut Self {
|
||||
self.only_files = Some(only_files);
|
||||
self
|
||||
|
|
@ -472,7 +489,7 @@ impl ManagedTorrentBuilder {
|
|||
self
|
||||
}
|
||||
|
||||
pub(crate) fn build(self, span: tracing::Span) -> anyhow::Result<ManagedTorrentHandle> {
|
||||
pub(crate) fn build(mut self, span: tracing::Span) -> anyhow::Result<ManagedTorrentHandle> {
|
||||
let lengths = Lengths::from_torrent(&self.info)?;
|
||||
let info = Arc::new(ManagedTorrentInfo {
|
||||
span,
|
||||
|
|
@ -499,6 +516,7 @@ impl ManagedTorrentBuilder {
|
|||
locked: RwLock::new(ManagedTorrentLocked {
|
||||
state: ManagedTorrentState::Initializing(initializing),
|
||||
}),
|
||||
cancellation_token: self.cancellation_token.take().unwrap_or_default(),
|
||||
info,
|
||||
}))
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue