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

Use tokio_util::CancellationToken everywhere

Browse source
This commit is contained in:
Igor Katson 2023-12-07 08:10:17 +00:00
parent 53868ad45e
commit bed7433d8e
No known key found for this signature in database
GPG key ID: B4EC22B66D61A3F5
16 changed files with 176 additions and 178 deletions
Download patch file Download diff file Expand all files Collapse all files

42
crates/dht/src/dht.rs
View file

@ -23,10 +23,14 @@ use anyhow::{bail, Context};
use backoff::{backoff::Backoff, ExponentialBackoffBuilder};
use bencode::ByteString;
use dashmap::DashMap;
use futures::{stream::FuturesUnordered, Future, Stream, StreamExt, TryFutureExt};
use futures::{stream::FuturesUnordered, Stream, StreamExt, TryFutureExt};
use leaky_bucket::RateLimiter;
use librqbit_core::{id20::Id20, peer_id::generate_peer_id, spawn_utils::spawn};
use librqbit_core::{
id20::Id20,
peer_id::generate_peer_id,
spawn_utils::{spawn, spawn_with_cancel},
};
use parking_lot::RwLock;
use serde::Serialize;
@ -35,6 +39,7 @@ use tokio::{
sync::mpsc::{channel, unbounded_channel, Sender, UnboundedReceiver, UnboundedSender},
};
use tokio_util::sync::CancellationToken;
use tracing::{debug, debug_span, error, error_span, info, trace, warn, Instrument};
#[derive(Debug, Serialize)]
@ -535,6 +540,8 @@ pub struct DhtState {
// This is to send raw messages
worker_sender: UnboundedSender<WorkerSendRequest>,
cancellation_token: CancellationToken,
pub(crate) peer_store: PeerStore,
}
@ -545,6 +552,7 @@ impl DhtState {
routing_table: Option<RoutingTable>,
listen_addr: SocketAddr,
peer_store: PeerStore,
cancellation_token: CancellationToken,
) -> Self {
let routing_table = routing_table.unwrap_or_else(|| RoutingTable::new(id, None));
Self {
@ -556,6 +564,7 @@ impl DhtState {
listen_addr,
rate_limiter: make_rate_limiter(),
peer_store,
cancellation_token,
}
}
@ -1124,21 +1133,18 @@ pub struct DhtConfig {
pub routing_table: Option<RoutingTable>,
pub listen_addr: Option<SocketAddr>,
pub peer_store: Option<PeerStore>,
pub cancellation_token: Option<CancellationToken>,
}
impl DhtState {
pub async fn new() -> anyhow::Result<(
Arc<Self>,
impl Future<Output = anyhow::Result<()>> + Send + Sync + 'static,
)> {
pub async fn new() -> anyhow::Result<Arc<Self>> {
Self::with_config(DhtConfig::default()).await
}
pub async fn with_config(
config: DhtConfig,
) -> anyhow::Result<(
Arc<Self>,
impl Future<Output = anyhow::Result<()>> + Send + Sync + 'static,
)> {
pub fn cancellation_token(&self) -> &CancellationToken {
&self.cancellation_token
}
pub async fn with_config(mut config: DhtConfig) -> anyhow::Result<Arc<Self>> {
let socket = match config.listen_addr {
Some(addr) => UdpSocket::bind(addr)
.await
@ -1159,6 +1165,8 @@ impl DhtState {
.bootstrap_addrs
.unwrap_or_else(|| crate::DHT_BOOTSTRAP.iter().map(|v| v.to_string()).collect());
let token = config.cancellation_token.take().unwrap_or_default();
let (in_tx, in_rx) = unbounded_channel();
let state = Arc::new(Self::new_internal(
peer_id,
@ -1166,17 +1174,17 @@ impl DhtState {
config.routing_table,
listen_addr,
config.peer_store.unwrap_or_else(|| PeerStore::new(peer_id)),
token,
));
let run_worker = {
spawn_with_cancel(error_span!("dht"), state.cancellation_token.clone(), {
let state = state.clone();
async move {
let worker = DhtWorker { socket, dht: state };
worker.start(in_rx, &bootstrap_addrs).await?;
Ok(())
worker.start(in_rx, &bootstrap_addrs).await
}
};
Ok((state, run_worker))
});
Ok(state)
}
pub fn get_peers(

13
crates/dht/src/lib.rs
View file

@ -10,7 +10,6 @@ use std::time::Duration;
pub use crate::dht::DhtStats;
pub use crate::dht::{DhtConfig, DhtState, RequestPeersStream};
use futures::Future;
pub use librqbit_core::id20::Id20;
pub use persistence::{PersistentDht, PersistentDhtConfig};
@ -27,19 +26,11 @@ pub struct DhtBuilder {}
impl DhtBuilder {
#[allow(clippy::new_ret_no_self)]
pub async fn new() -> anyhow::Result<(
Dht,
impl Future<Output = anyhow::Result<()>> + Send + Sync + 'static,
)> {
pub async fn new() -> anyhow::Result<Dht> {
DhtState::new().await
}
pub async fn with_config(
config: DhtConfig,
) -> anyhow::Result<(
Dht,
impl Future<Output = anyhow::Result<()>> + Send + Sync + 'static,
)> {
pub async fn with_config(config: DhtConfig) -> anyhow::Result<Dht> {
DhtState::with_config(config).await
}
}

64
crates/dht/src/persistence.rs
View file

@ -1,16 +1,17 @@
// TODO: this now stores only the routing table, but we also need AT LEAST the same socket address...
use futures::Future;
use librqbit_core::directories::get_configuration_directory;
use librqbit_core::spawn_utils::spawn_with_cancel;
use serde::{Deserialize, Serialize};
use std::fs::OpenOptions;
use std::io::{BufReader, BufWriter};
use std::net::SocketAddr;
use std::path::{Path, PathBuf};
use std::time::Duration;
use tokio_util::sync::CancellationToken;
use anyhow::Context;
use tracing::{debug, error, info, trace, warn};
use tracing::{debug, error, error_span, info, trace, warn};
use crate::peer_store::PeerStore;
use crate::routing_table::RoutingTable;
@ -76,11 +77,8 @@ impl PersistentDht {
pub async fn create(
config: Option<PersistentDhtConfig>,
) -> anyhow::Result<(
Dht,
impl Future<Output = anyhow::Result<()>> + Send + Sync + 'static,
impl Future<Output = anyhow::Result<()>> + Send + Sync + 'static,
)> {
cancellation_token: Option<CancellationToken>,
) -> anyhow::Result<Dht> {
let mut config = config.unwrap_or_default();
let config_filename = match config.config_filename.take() {
Some(config_filename) => config_filename,
@ -129,35 +127,41 @@ impl PersistentDht {
routing_table,
listen_addr,
peer_store,
cancellation_token,
..Default::default()
};
let (dht, run_worker) = DhtState::with_config(dht_config).await?;
let dht = DhtState::with_config(dht_config).await?;
spawn_with_cancel(
error_span!("dht_persistence"),
dht.cancellation_token().clone(),
{
let dht = dht.clone();
let dump_interval = config
.dump_interval
.unwrap_or_else(|| Duration::from_secs(3));
async move {
let tempfile_name = {
let file_name = format!("dht.json.tmp.{}", std::process::id());
let mut tmp = config_filename.clone();
tmp.set_file_name(file_name);
tmp
};
let run_persistence = {
let dht = dht.clone();
let dump_interval = config
.dump_interval
.unwrap_or_else(|| Duration::from_secs(3));
async move {
let tempfile_name = {
let file_name = format!("dht.json.tmp.{}", std::process::id());
let mut tmp = config_filename.clone();
tmp.set_file_name(file_name);
tmp
};
loop {
trace!("sleeping for {:?}", &dump_interval);
tokio::time::sleep(dump_interval).await;
loop {
trace!("sleeping for {:?}", &dump_interval);
tokio::time::sleep(dump_interval).await;
match dump_dht(&dht, &config_filename, &tempfile_name) {
Ok(_) => debug!("dumped DHT to {:?}", &config_filename),
Err(e) => error!("error dumping DHT to {:?}: {:#}", &config_filename, e),
match dump_dht(&dht, &config_filename, &tempfile_name) {
Ok(_) => debug!("dumped DHT to {:?}", &config_filename),
Err(e) => {
error!("error dumping DHT to {:?}: {:#}", &config_filename, e)
}
}
}
}
}
};
},
);
Ok((dht, run_worker, run_persistence))
Ok(dht)
}
}