YGGverse/aquatic
1
0
Fork 0
mirror of https://github.com/YGGverse/aquatic.git synced 2026-03-31 17:55:36 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

http: send response sender to request workers

Browse source
This commit is contained in:
Joakim Frostegård 2022-04-10 02:30:49 +02:00
parent 02b9c9d290
commit f083cac9e0
6 changed files with 78 additions and 221 deletions
Download patch file Download diff file Expand all files Collapse all files

1
Cargo.lock generated
View file

@ -90,6 +90,7 @@ dependencies = [
"aquatic_toml_config", "aquatic_toml_config",
"cfg-if", "cfg-if",
"either", "either",
"futures",
"futures-lite", "futures-lite",
"futures-rustls", "futures-rustls",
"glommio", "glommio",

1
aquatic_http/Cargo.toml
View file

@ -23,6 +23,7 @@ aquatic_toml_config = { version = "0.2.0", path = "../aquatic_toml_config" }
anyhow = "1" anyhow = "1"
cfg-if = "1" cfg-if = "1"
either = "1" either = "1"
futures = "0.3"
futures-lite = "1" futures-lite = "1"
futures-rustls = "0.22" futures-rustls = "0.22"
glommio = "0.7" glommio = "0.7"

36
aquatic_http/src/common.rs
View file

@ -9,6 +9,7 @@ use aquatic_http_protocol::{
request::{AnnounceRequest, ScrapeRequest}, request::{AnnounceRequest, ScrapeRequest},
response::{AnnounceResponse, ScrapeResponse}, response::{AnnounceResponse, ScrapeResponse},
}; };
use glommio::channels::shared_channel::SharedSender;
#[derive(Copy, Clone, Debug)] #[derive(Copy, Clone, Debug)]
pub struct ConsumerId(pub usize); pub struct ConsumerId(pub usize);
@ -21,46 +22,15 @@ pub enum ChannelRequest {
Announce { Announce {
request: AnnounceRequest, request: AnnounceRequest,
peer_addr: CanonicalSocketAddr, peer_addr: CanonicalSocketAddr,
connection_id: ConnectionId, response_sender: SharedSender<AnnounceResponse>,
response_consumer_id: ConsumerId,
}, },
Scrape { Scrape {
request: ScrapeRequest, request: ScrapeRequest,
peer_addr: CanonicalSocketAddr, peer_addr: CanonicalSocketAddr,
connection_id: ConnectionId, response_sender: SharedSender<ScrapeResponse>,
response_consumer_id: ConsumerId,
}, },
} }
#[derive(Debug)]
pub enum ChannelResponse {
Announce {
response: AnnounceResponse,
peer_addr: CanonicalSocketAddr,
connection_id: ConnectionId,
},
Scrape {
response: ScrapeResponse,
peer_addr: CanonicalSocketAddr,
connection_id: ConnectionId,
},
}
impl ChannelResponse {
pub fn get_connection_id(&self) -> ConnectionId {
match self {
Self::Announce { connection_id, .. } => *connection_id,
Self::Scrape { connection_id, .. } => *connection_id,
}
}
pub fn get_peer_addr(&self) -> CanonicalSocketAddr {
match self {
Self::Announce { peer_addr, .. } => *peer_addr,
Self::Scrape { peer_addr, .. } => *peer_addr,
}
}
}
#[derive(Default, Clone)] #[derive(Default, Clone)]
pub struct State { pub struct State {
pub access_list: Arc<AccessListArcSwap>, pub access_list: Arc<AccessListArcSwap>,

14
aquatic_http/src/lib.rs
View file

@ -37,7 +37,6 @@ pub fn run(config: Config) -> ::anyhow::Result<()> {
let num_peers = config.socket_workers + config.request_workers; let num_peers = config.socket_workers + config.request_workers;
let request_mesh_builder = MeshBuilder::partial(num_peers, SHARED_CHANNEL_SIZE); let request_mesh_builder = MeshBuilder::partial(num_peers, SHARED_CHANNEL_SIZE);
let response_mesh_builder = MeshBuilder::partial(num_peers, SHARED_CHANNEL_SIZE);
let (sentinel_watcher, sentinel) = PanicSentinelWatcher::create_with_sentinel(); let (sentinel_watcher, sentinel) = PanicSentinelWatcher::create_with_sentinel();
let priv_dropper = PrivilegeDropper::new(config.privileges.clone(), config.socket_workers); let priv_dropper = PrivilegeDropper::new(config.privileges.clone(), config.socket_workers);
@ -55,7 +54,6 @@ pub fn run(config: Config) -> ::anyhow::Result<()> {
let state = state.clone(); let state = state.clone();
let tls_config = tls_config.clone(); let tls_config = tls_config.clone();
let request_mesh_builder = request_mesh_builder.clone(); let request_mesh_builder = request_mesh_builder.clone();
let response_mesh_builder = response_mesh_builder.clone();
let priv_dropper = priv_dropper.clone(); let priv_dropper = priv_dropper.clone();
let placement = get_worker_placement( let placement = get_worker_placement(
@ -74,7 +72,6 @@ pub fn run(config: Config) -> ::anyhow::Result<()> {
state, state,
tls_config, tls_config,
request_mesh_builder, request_mesh_builder,
response_mesh_builder,
priv_dropper, priv_dropper,
) )
.await .await
@ -89,7 +86,6 @@ pub fn run(config: Config) -> ::anyhow::Result<()> {
let config = config.clone(); let config = config.clone();
let state = state.clone(); let state = state.clone();
let request_mesh_builder = request_mesh_builder.clone(); let request_mesh_builder = request_mesh_builder.clone();
let response_mesh_builder = response_mesh_builder.clone();
let placement = get_worker_placement( let placement = get_worker_placement(
&config.cpu_pinning, &config.cpu_pinning,
@ -101,14 +97,8 @@ pub fn run(config: Config) -> ::anyhow::Result<()> {
let executor = builder let executor = builder
.spawn(move || async move { .spawn(move || async move {
workers::request::run_request_worker( workers::request::run_request_worker(sentinel, config, state, request_mesh_builder)
sentinel, .await
config,
state,
request_mesh_builder,
response_mesh_builder,
)
.await
}) })
.map_err(|err| anyhow::anyhow!("Spawning executor failed: {:#}", err))?; .map_err(|err| anyhow::anyhow!("Spawning executor failed: {:#}", err))?;

118
aquatic_http/src/workers/request.rs
View file

@ -8,7 +8,7 @@ use std::time::Instant;
use either::Either; use either::Either;
use futures_lite::{Stream, StreamExt}; use futures_lite::{Stream, StreamExt};
use glommio::channels::channel_mesh::{MeshBuilder, Partial, Role, Senders}; use glommio::channels::channel_mesh::{MeshBuilder, Partial, Role};
use glommio::timer::TimerActionRepeat; use glommio::timer::TimerActionRepeat;
use glommio::{enclose, prelude::*}; use glommio::{enclose, prelude::*};
use rand::prelude::SmallRng; use rand::prelude::SmallRng;
@ -33,23 +33,6 @@ pub trait Ip: ::std::fmt::Debug + Copy + Eq + ::std::hash::Hash {}
impl Ip for Ipv4Addr {} impl Ip for Ipv4Addr {}
impl Ip for Ipv6Addr {} impl Ip for Ipv6Addr {}
#[derive(Clone, Copy, Debug)]
pub struct ConnectionMeta {
/// Index of socket worker responsible for this connection. Required for
/// sending back response through correct channel to correct worker.
pub response_consumer_id: ConsumerId,
pub peer_addr: CanonicalSocketAddr,
/// Connection id local to socket worker
pub connection_id: ConnectionId,
}
#[derive(Clone, Copy, Debug)]
pub struct PeerConnectionMeta<I: Ip> {
pub response_consumer_id: ConsumerId,
pub connection_id: ConnectionId,
pub peer_ip_address: I,
}
#[derive(PartialEq, Eq, Clone, Copy, Debug)] #[derive(PartialEq, Eq, Clone, Copy, Debug)]
pub enum PeerStatus { pub enum PeerStatus {
Seeding, Seeding,
@ -75,7 +58,7 @@ impl PeerStatus {
#[derive(Debug, Clone, Copy)] #[derive(Debug, Clone, Copy)]
pub struct Peer<I: Ip> { pub struct Peer<I: Ip> {
pub connection_meta: PeerConnectionMeta<I>, pub ip_address: I,
pub port: u16, pub port: u16,
pub status: PeerStatus, pub status: PeerStatus,
pub valid_until: ValidUntil, pub valid_until: ValidUntil,
@ -84,7 +67,7 @@ pub struct Peer<I: Ip> {
impl<I: Ip> Peer<I> { impl<I: Ip> Peer<I> {
pub fn to_response_peer(&self) -> ResponsePeer<I> { pub fn to_response_peer(&self) -> ResponsePeer<I> {
ResponsePeer { ResponsePeer {
ip_address: self.connection_meta.peer_ip_address, ip_address: self.ip_address,
port: self.port, port: self.port,
} }
} }
@ -179,12 +162,8 @@ pub async fn run_request_worker(
config: Config, config: Config,
state: State, state: State,
request_mesh_builder: MeshBuilder<ChannelRequest, Partial>, request_mesh_builder: MeshBuilder<ChannelRequest, Partial>,
response_mesh_builder: MeshBuilder<ChannelResponse, Partial>,
) { ) {
let (_, mut request_receivers) = request_mesh_builder.join(Role::Consumer).await.unwrap(); let (_, mut request_receivers) = request_mesh_builder.join(Role::Consumer).await.unwrap();
let (response_senders, _) = response_mesh_builder.join(Role::Producer).await.unwrap();
let response_senders = Rc::new(response_senders);
let torrents = Rc::new(RefCell::new(TorrentMaps::default())); let torrents = Rc::new(RefCell::new(TorrentMaps::default()));
let access_list = state.access_list; let access_list = state.access_list;
@ -204,7 +183,6 @@ pub async fn run_request_worker(
let handle = spawn_local(handle_request_stream( let handle = spawn_local(handle_request_stream(
config.clone(), config.clone(),
torrents.clone(), torrents.clone(),
response_senders.clone(),
receiver, receiver,
)) ))
.detach(); .detach();
@ -217,12 +195,8 @@ pub async fn run_request_worker(
} }
} }
async fn handle_request_stream<S>( async fn handle_request_stream<S>(config: Config, torrents: Rc<RefCell<TorrentMaps>>, mut stream: S)
config: Config, where
torrents: Rc<RefCell<TorrentMaps>>,
response_senders: Rc<Senders<ChannelResponse>>,
mut stream: S,
) where
S: Stream<Item = ChannelRequest> + ::std::marker::Unpin, S: Stream<Item = ChannelRequest> + ::std::marker::Unpin,
{ {
let mut rng = SmallRng::from_entropy(); let mut rng = SmallRng::from_entropy();
@ -239,68 +213,38 @@ async fn handle_request_stream<S>(
})); }));
while let Some(channel_request) = stream.next().await { while let Some(channel_request) = stream.next().await {
let (response, consumer_id) = match channel_request { match channel_request {
ChannelRequest::Announce { ChannelRequest::Announce {
request, request,
peer_addr, peer_addr,
response_consumer_id, response_sender,
connection_id,
} => { } => {
let meta = ConnectionMeta {
response_consumer_id,
connection_id,
peer_addr,
};
let response = handle_announce_request( let response = handle_announce_request(
&config, &config,
&mut rng, &mut rng,
&mut torrents.borrow_mut(), &mut torrents.borrow_mut(),
peer_valid_until.borrow().to_owned(), peer_valid_until.borrow().to_owned(),
meta, peer_addr,
request, request,
); );
let response = ChannelResponse::Announce { if let Err(err) = response_sender.connect().await.send(response).await {
response, ::log::error!("request worker could not send announce response: {:#}", err);
peer_addr, }
connection_id,
};
(response, response_consumer_id)
} }
ChannelRequest::Scrape { ChannelRequest::Scrape {
request, request,
peer_addr, peer_addr,
response_consumer_id, response_sender,
connection_id,
} => { } => {
let meta = ConnectionMeta {
response_consumer_id,
connection_id,
peer_addr,
};
let response = let response =
handle_scrape_request(&config, &mut torrents.borrow_mut(), meta, request); handle_scrape_request(&config, &mut torrents.borrow_mut(), peer_addr, request);
let response = ChannelResponse::Scrape { if let Err(err) = response_sender.connect().await.send(response).await {
response, ::log::error!("request worker could not send scrape response: {:#}", err);
peer_addr, }
connection_id,
};
(response, response_consumer_id)
} }
}; };
::log::debug!("preparing to send response to channel: {:?}", response);
if let Err(err) = response_senders.try_send_to(consumer_id.0, response) {
::log::warn!("response_sender.try_send: {:?}", err);
}
yield_if_needed().await;
} }
} }
@ -309,24 +253,18 @@ pub fn handle_announce_request(
rng: &mut impl Rng, rng: &mut impl Rng,
torrent_maps: &mut TorrentMaps, torrent_maps: &mut TorrentMaps,
valid_until: ValidUntil, valid_until: ValidUntil,
meta: ConnectionMeta, peer_addr: CanonicalSocketAddr,
request: AnnounceRequest, request: AnnounceRequest,
) -> AnnounceResponse { ) -> AnnounceResponse {
match meta.peer_addr.get().ip() { match peer_addr.get().ip() {
IpAddr::V4(peer_ip_address) => { IpAddr::V4(peer_ip_address) => {
let torrent_data: &mut TorrentData<Ipv4Addr> = let torrent_data: &mut TorrentData<Ipv4Addr> =
torrent_maps.ipv4.entry(request.info_hash).or_default(); torrent_maps.ipv4.entry(request.info_hash).or_default();
let peer_connection_meta = PeerConnectionMeta {
response_consumer_id: meta.response_consumer_id,
connection_id: meta.connection_id,
peer_ip_address,
};
let (seeders, leechers, response_peers) = upsert_peer_and_get_response_peers( let (seeders, leechers, response_peers) = upsert_peer_and_get_response_peers(
config, config,
rng, rng,
peer_connection_meta, peer_ip_address,
torrent_data, torrent_data,
request, request,
valid_until, valid_until,
@ -347,16 +285,10 @@ pub fn handle_announce_request(
let torrent_data: &mut TorrentData<Ipv6Addr> = let torrent_data: &mut TorrentData<Ipv6Addr> =
torrent_maps.ipv6.entry(request.info_hash).or_default(); torrent_maps.ipv6.entry(request.info_hash).or_default();
let peer_connection_meta = PeerConnectionMeta {
response_consumer_id: meta.response_consumer_id,
connection_id: meta.connection_id,
peer_ip_address,
};
let (seeders, leechers, response_peers) = upsert_peer_and_get_response_peers( let (seeders, leechers, response_peers) = upsert_peer_and_get_response_peers(
config, config,
rng, rng,
peer_connection_meta, peer_ip_address,
torrent_data, torrent_data,
request, request,
valid_until, valid_until,
@ -380,7 +312,7 @@ pub fn handle_announce_request(
pub fn upsert_peer_and_get_response_peers<I: Ip>( pub fn upsert_peer_and_get_response_peers<I: Ip>(
config: &Config, config: &Config,
rng: &mut impl Rng, rng: &mut impl Rng,
request_sender_meta: PeerConnectionMeta<I>, peer_ip_address: I,
torrent_data: &mut TorrentData<I>, torrent_data: &mut TorrentData<I>,
request: AnnounceRequest, request: AnnounceRequest,
valid_until: ValidUntil, valid_until: ValidUntil,
@ -391,7 +323,7 @@ pub fn upsert_peer_and_get_response_peers<I: Ip>(
PeerStatus::from_event_and_bytes_left(request.event, Some(request.bytes_left)); PeerStatus::from_event_and_bytes_left(request.event, Some(request.bytes_left));
let peer = Peer { let peer = Peer {
connection_meta: request_sender_meta, ip_address: peer_ip_address,
port: request.port, port: request.port,
status: peer_status, status: peer_status,
valid_until, valid_until,
@ -402,7 +334,7 @@ pub fn upsert_peer_and_get_response_peers<I: Ip>(
let ip_or_key = request let ip_or_key = request
.key .key
.map(Either::Right) .map(Either::Right)
.unwrap_or_else(|| Either::Left(request_sender_meta.peer_ip_address)); .unwrap_or_else(|| Either::Left(peer_ip_address));
let peer_map_key = PeerMapKey { let peer_map_key = PeerMapKey {
peer_id: request.peer_id, peer_id: request.peer_id,
@ -462,7 +394,7 @@ pub fn upsert_peer_and_get_response_peers<I: Ip>(
pub fn handle_scrape_request( pub fn handle_scrape_request(
config: &Config, config: &Config,
torrent_maps: &mut TorrentMaps, torrent_maps: &mut TorrentMaps,
meta: ConnectionMeta, peer_addr: CanonicalSocketAddr,
request: ScrapeRequest, request: ScrapeRequest,
) -> ScrapeResponse { ) -> ScrapeResponse {
let num_to_take = request let num_to_take = request
@ -474,7 +406,7 @@ pub fn handle_scrape_request(
files: BTreeMap::new(), files: BTreeMap::new(),
}; };
let peer_ip = meta.peer_addr.get().ip(); let peer_ip = peer_addr.get().ip();
// If request.info_hashes is empty, don't return scrape for all // If request.info_hashes is empty, don't return scrape for all
// torrents, even though reference server does it. It is too expensive. // torrents, even though reference server does it. It is too expensive.

129
aquatic_http/src/workers/socket.rs
View file

@ -16,12 +16,12 @@ use aquatic_http_protocol::response::{
FailureResponse, Response, ScrapeResponse, ScrapeStatistics, FailureResponse, Response, ScrapeResponse, ScrapeStatistics,
}; };
use either::Either; use either::Either;
use futures::stream::FuturesUnordered;
use futures_lite::{AsyncReadExt, AsyncWriteExt, StreamExt}; use futures_lite::{AsyncReadExt, AsyncWriteExt, StreamExt};
use futures_rustls::server::TlsStream; use futures_rustls::server::TlsStream;
use futures_rustls::TlsAcceptor; use futures_rustls::TlsAcceptor;
use glommio::channels::channel_mesh::{MeshBuilder, Partial, Role, Senders}; use glommio::channels::channel_mesh::{MeshBuilder, Partial, Role, Senders};
use glommio::channels::local_channel::{new_bounded, LocalReceiver, LocalSender}; use glommio::channels::shared_channel::{self, SharedReceiver};
use glommio::channels::shared_channel::ConnectedReceiver;
use glommio::net::{TcpListener, TcpStream}; use glommio::net::{TcpListener, TcpStream};
use glommio::task::JoinHandle; use glommio::task::JoinHandle;
use glommio::timer::TimerActionRepeat; use glommio::timer::TimerActionRepeat;
@ -49,7 +49,6 @@ struct PendingScrapeResponse {
struct ConnectionReference { struct ConnectionReference {
task_handle: Option<JoinHandle<()>>, task_handle: Option<JoinHandle<()>>,
response_sender: LocalSender<ChannelResponse>,
valid_until: ValidUntil, valid_until: ValidUntil,
} }
@ -59,7 +58,6 @@ pub async fn run_socket_worker(
state: State, state: State,
tls_config: Arc<RustlsConfig>, tls_config: Arc<RustlsConfig>,
request_mesh_builder: MeshBuilder<ChannelRequest, Partial>, request_mesh_builder: MeshBuilder<ChannelRequest, Partial>,
response_mesh_builder: MeshBuilder<ChannelResponse, Partial>,
priv_dropper: PrivilegeDropper, priv_dropper: PrivilegeDropper,
) { ) {
let config = Rc::new(config); let config = Rc::new(config);
@ -70,9 +68,6 @@ pub async fn run_socket_worker(
let (request_senders, _) = request_mesh_builder.join(Role::Producer).await.unwrap(); let (request_senders, _) = request_mesh_builder.join(Role::Producer).await.unwrap();
let request_senders = Rc::new(request_senders); let request_senders = Rc::new(request_senders);
let (_, mut response_receivers) = response_mesh_builder.join(Role::Consumer).await.unwrap();
let response_consumer_id = ConsumerId(response_receivers.consumer_id().unwrap());
let connection_slab = Rc::new(RefCell::new(Slab::new())); let connection_slab = Rc::new(RefCell::new(Slab::new()));
TimerActionRepeat::repeat(enclose!((config, connection_slab) move || { TimerActionRepeat::repeat(enclose!((config, connection_slab) move || {
@ -82,24 +77,13 @@ pub async fn run_socket_worker(
) )
})); }));
for (_, response_receiver) in response_receivers.streams() {
spawn_local(receive_responses(
response_receiver,
connection_slab.clone(),
))
.detach();
}
let mut incoming = listener.incoming(); let mut incoming = listener.incoming();
while let Some(stream) = incoming.next().await { while let Some(stream) = incoming.next().await {
match stream { match stream {
Ok(stream) => { Ok(stream) => {
let (response_sender, response_receiver) = new_bounded(config.request_workers);
let key = connection_slab.borrow_mut().insert(ConnectionReference { let key = connection_slab.borrow_mut().insert(ConnectionReference {
task_handle: None, task_handle: None,
response_sender,
valid_until: ValidUntil::new(config.cleaning.max_connection_idle), valid_until: ValidUntil::new(config.cleaning.max_connection_idle),
}); });
@ -108,8 +92,6 @@ pub async fn run_socket_worker(
config, config,
access_list, access_list,
request_senders, request_senders,
response_receiver,
response_consumer_id,
ConnectionId(key), ConnectionId(key),
tls_config, tls_config,
connection_slab.clone(), connection_slab.clone(),
@ -158,28 +140,10 @@ async fn clean_connections(
)) ))
} }
async fn receive_responses(
mut response_receiver: ConnectedReceiver<ChannelResponse>,
connection_references: Rc<RefCell<Slab<ConnectionReference>>>,
) {
while let Some(channel_response) = response_receiver.next().await {
if let Some(reference) = connection_references
.borrow()
.get(channel_response.get_connection_id().0)
{
if let Err(err) = reference.response_sender.try_send(channel_response) {
::log::error!("Couldn't send response to local receiver: {:?}", err);
}
}
}
}
struct Connection { struct Connection {
config: Rc<Config>, config: Rc<Config>,
access_list_cache: AccessListCache, access_list_cache: AccessListCache,
request_senders: Rc<Senders<ChannelRequest>>, request_senders: Rc<Senders<ChannelRequest>>,
response_receiver: LocalReceiver<ChannelResponse>,
response_consumer_id: ConsumerId,
connection_slab: Rc<RefCell<Slab<ConnectionReference>>>, connection_slab: Rc<RefCell<Slab<ConnectionReference>>>,
stream: TlsStream<TcpStream>, stream: TlsStream<TcpStream>,
peer_addr: CanonicalSocketAddr, peer_addr: CanonicalSocketAddr,
@ -194,8 +158,6 @@ impl Connection {
config: Rc<Config>, config: Rc<Config>,
access_list: Arc<AccessListArcSwap>, access_list: Arc<AccessListArcSwap>,
request_senders: Rc<Senders<ChannelRequest>>, request_senders: Rc<Senders<ChannelRequest>>,
response_receiver: LocalReceiver<ChannelResponse>,
response_consumer_id: ConsumerId,
connection_id: ConnectionId, connection_id: ConnectionId,
tls_config: Arc<RustlsConfig>, tls_config: Arc<RustlsConfig>,
connection_slab: Rc<RefCell<Slab<ConnectionReference>>>, connection_slab: Rc<RefCell<Slab<ConnectionReference>>>,
@ -217,8 +179,6 @@ impl Connection {
config: config.clone(), config: config.clone(),
access_list_cache: create_access_list_cache(&access_list), access_list_cache: create_access_list_cache(&access_list),
request_senders: request_senders.clone(), request_senders: request_senders.clone(),
response_receiver,
response_consumer_id,
connection_slab, connection_slab,
stream, stream,
peer_addr, peer_addr,
@ -324,11 +284,12 @@ impl Connection {
.load() .load()
.allows(self.config.access_list.mode, &info_hash.0) .allows(self.config.access_list.mode, &info_hash.0)
{ {
let (response_sender, response_receiver) = shared_channel::new_bounded(1);
let request = ChannelRequest::Announce { let request = ChannelRequest::Announce {
request, request,
connection_id: self.connection_id,
response_consumer_id: self.response_consumer_id,
peer_addr: self.peer_addr, peer_addr: self.peer_addr,
response_sender,
}; };
let consumer_index = calculate_request_consumer_index(&self.config, info_hash); let consumer_index = calculate_request_consumer_index(&self.config, info_hash);
@ -339,7 +300,13 @@ impl Connection {
.await .await
.unwrap(); .unwrap();
self.wait_for_response(None).await response_receiver
.connect()
.await
.recv()
.await
.ok_or_else(|| anyhow::anyhow!("response sender closed"))
.map(Response::Announce)
} else { } else {
let response = Response::Failure(FailureResponse { let response = Response::Failure(FailureResponse {
failure_reason: "Info hash not allowed".into(), failure_reason: "Info hash not allowed".into(),
@ -360,13 +327,17 @@ impl Connection {
} }
let pending_worker_responses = info_hashes_by_worker.len(); let pending_worker_responses = info_hashes_by_worker.len();
let mut response_receivers = Vec::new();
for (consumer_index, info_hashes) in info_hashes_by_worker { for (consumer_index, info_hashes) in info_hashes_by_worker {
let (response_sender, response_receiver) = shared_channel::new_bounded(1);
response_receivers.push(response_receiver);
let request = ChannelRequest::Scrape { let request = ChannelRequest::Scrape {
request: ScrapeRequest { info_hashes }, request: ScrapeRequest { info_hashes },
peer_addr: self.peer_addr, peer_addr: self.peer_addr,
response_consumer_id: self.response_consumer_id, response_sender,
connection_id: self.connection_id,
}; };
// Only fails when receiver is closed // Only fails when receiver is closed
@ -381,53 +352,45 @@ impl Connection {
stats: Default::default(), stats: Default::default(),
}; };
self.wait_for_response(Some(pending_scrape_response)).await self.wait_for_scrape_responses(response_receivers, pending_scrape_response)
.await
} }
} }
} }
/// Wait for announce response or partial scrape responses to arrive, /// Wait for partial scrape responses to arrive,
/// return full response /// return full response
async fn wait_for_response( async fn wait_for_scrape_responses(
&self, &self,
mut opt_pending_scrape_response: Option<PendingScrapeResponse>, response_receivers: Vec<SharedReceiver<ScrapeResponse>>,
mut pending: PendingScrapeResponse,
) -> anyhow::Result<Response> { ) -> anyhow::Result<Response> {
let mut responses = response_receivers
.into_iter()
.map(|receiver| async { receiver.connect().await.recv().await })
.collect::<FuturesUnordered<_>>();
loop { loop {
let channel_response = self let response = responses
.response_receiver .next()
.recv()
.await .await
.expect("wait_for_response: can't receive response, sender is closed"); .ok_or_else(|| anyhow::anyhow!("wait_for_scrape_response: all messages received"))?
.ok_or_else(|| {
anyhow::anyhow!(
"wait_for_scrape_response: can't receive response, sender is closed"
)
})?;
if channel_response.get_peer_addr() != self.peer_addr { pending.stats.extend(response.files);
return Err(anyhow::anyhow!("peer addresses didn't match")); pending.pending_worker_responses -= 1;
if pending.pending_worker_responses == 0 {
let response = Response::Scrape(ScrapeResponse {
files: pending.stats,
});
break Ok(response);
} }
match channel_response {
ChannelResponse::Announce { response, .. } => {
break Ok(Response::Announce(response));
}
ChannelResponse::Scrape { response, .. } => {
if let Some(mut pending) = opt_pending_scrape_response.take() {
pending.stats.extend(response.files);
pending.pending_worker_responses -= 1;
if pending.pending_worker_responses == 0 {
let response = Response::Scrape(ScrapeResponse {
files: pending.stats,
});
break Ok(response);
} else {
opt_pending_scrape_response = Some(pending);
}
} else {
return Err(anyhow::anyhow!(
"received channel scrape response without pending scrape response"
));
}
}
};
} }
} }