This commit is contained in:
Igor Katson 2021-07-13 16:28:53 +01:00
parent ace11186ef
commit 1b5e565aff
2 changed files with 27 additions and 28 deletions

View file

@ -44,6 +44,12 @@ struct DhtState {
next_transaction_id: u16, next_transaction_id: u16,
outstanding_requests: Vec<OutstandingRequest>, outstanding_requests: Vec<OutstandingRequest>,
routing_table: RoutingTable, routing_table: RoutingTable,
// This sender sends requests to the worker.
// It is unbounded so that the methods on Dht state don't need to be async.
// If the methods on Dht state were async, we would have a problem, as it's behind
// a lock.
// Alternatively, we can lock only the parts that change, and use that internally inside DhtState...
sender: UnboundedSender<(Message<ByteString>, SocketAddr)>, sender: UnboundedSender<(Message<ByteString>, SocketAddr)>,
seen_peers: HashMap<Id20, HashSet<SocketAddr>>, seen_peers: HashMap<Id20, HashSet<SocketAddr>>,
@ -420,7 +426,6 @@ impl DhtWorker {
self, self,
in_tx: UnboundedSender<(Message<ByteString>, SocketAddr)>, in_tx: UnboundedSender<(Message<ByteString>, SocketAddr)>,
in_rx: UnboundedReceiver<(Message<ByteString>, SocketAddr)>, in_rx: UnboundedReceiver<(Message<ByteString>, SocketAddr)>,
mut request_rx: Receiver<(Request, UnboundedSender<Response>)>,
bootstrap_addrs: &[String], bootstrap_addrs: &[String],
) -> anyhow::Result<()> { ) -> anyhow::Result<()> {
let (out_tx, mut out_rx) = channel(1); let (out_tx, mut out_rx) = channel(1);
@ -501,7 +506,6 @@ impl Dht {
Self::with_bootstrap_addrs(DHT_BOOTSTRAP).await Self::with_bootstrap_addrs(DHT_BOOTSTRAP).await
} }
pub async fn with_bootstrap_addrs(bootstrap_addrs: &[&str]) -> anyhow::Result<Self> { pub async fn with_bootstrap_addrs(bootstrap_addrs: &[&str]) -> anyhow::Result<Self> {
let (request_tx, request_rx) = channel(1);
let socket = UdpSocket::bind("0.0.0.0:0") let socket = UdpSocket::bind("0.0.0.0:0")
.await .await
.context("error binding socket")?; .context("error binding socket")?;
@ -523,9 +527,7 @@ impl Dht {
peer_id, peer_id,
state, state,
}; };
let result = worker let result = worker.start(in_tx, in_rx, &bootstrap_addrs).await;
.start(in_tx, in_rx, request_rx, &bootstrap_addrs)
.await;
warn!("DHT worker finished with {:?}", result); warn!("DHT worker finished with {:?}", result);
} }
}); });

View file

@ -1,9 +1,9 @@
use std::{fs::File, io::Read, net::SocketAddr, str::FromStr, time::Duration}; use std::{fs::File, io::Read, net::SocketAddr, pin::Pin, str::FromStr, time::Duration};
use anyhow::Context; use anyhow::Context;
use clap::Clap; use clap::Clap;
use dht::{Dht, Id20}; use dht::{Dht, Id20};
use futures::StreamExt; use futures::{Stream, StreamExt};
use librqbit::{ use librqbit::{
dht_utils::{read_metainfo_from_peer_receiver, ReadMetainfoResult}, dht_utils::{read_metainfo_from_peer_receiver, ReadMetainfoResult},
generate_peer_id, generate_peer_id,
@ -209,15 +209,7 @@ async fn async_main(opts: Opts, spawner: BlockingSpawner) -> anyhow::Result<()>
.ok_or_else(|| anyhow::anyhow!("magnet links without DHT are not supported"))? .ok_or_else(|| anyhow::anyhow!("magnet links without DHT are not supported"))?
.get_peers(info_hash) .get_peers(info_hash)
.await?; .await?;
let dht_rx = Box::pin(dht_rx.filter_map(|addr| async move { let dht_rx = flatten_dht_peers_stream(dht_rx);
match addr {
Ok(addr) => Some(addr),
Err(e) => {
warn!("DHT peer receiver got an error: {:#}", e);
None
}
}
}));
let trackers = trackers let trackers = trackers
.into_iter() .into_iter()
@ -259,18 +251,8 @@ async fn async_main(opts: Opts, spawner: BlockingSpawner) -> anyhow::Result<()>
torrent_from_file(&opts.torrent_path)? torrent_from_file(&opts.torrent_path)?
}; };
let dht_rx = match dht { let dht_rx = match dht {
Some(dht) => Some(Box::pin( Some(dht) => Some(flatten_dht_peers_stream(
dht.get_peers(torrent.info_hash) dht.get_peers(torrent.info_hash).await?,
.await?
.filter_map(|r| async move {
match r {
Ok(addr) => Some(addr),
Err(e) => {
warn!("DHT peer receiver got an error: {:#}", e);
None
}
}
}),
)), )),
None => None, None => None,
}; };
@ -307,6 +289,21 @@ async fn async_main(opts: Opts, spawner: BlockingSpawner) -> anyhow::Result<()>
} }
} }
fn flatten_dht_peers_stream(
rx: impl Stream<Item = Result<SocketAddr, anyhow::Error>> + Unpin,
) -> impl Stream<Item = SocketAddr> + Unpin {
let rx = rx.filter_map(|addr| async move {
match addr {
Ok(addr) => Some(addr),
Err(e) => {
warn!("DHT peer receiver got an error: {:#}", e);
None
}
}
});
Box::pin(rx)
}
#[allow(clippy::too_many_arguments)] #[allow(clippy::too_many_arguments)]
async fn main_torrent_info( async fn main_torrent_info(
opts: Opts, opts: Opts,