mirror of
https://github.com/YGGverse/aquatic.git
synced 2026-03-31 17:55:36 +00:00
ws: improve socket worker connection code
This commit is contained in:
Joakim Frostegård
parent
fe5ccf6646
commit
7b2a7a4f46
1 changed files with 229 additions and 242 deletions
|
|
@ -229,11 +229,14 @@ impl<S: futures::AsyncRead + futures::AsyncWrite + Unpin> ConnectionReader<S> {
|
|||
match &message {
|
||||
tungstenite::Message::Text(_) | tungstenite::Message::Binary(_) => {
|
||||
match InMessage::from_ws_message(message) {
|
||||
Ok(in_message) => {
|
||||
self.handle_in_message(in_message).await?;
|
||||
Ok(InMessage::AnnounceRequest(request)) => {
|
||||
self.handle_announce_request(request).await?;
|
||||
}
|
||||
Ok(InMessage::ScrapeRequest(request)) => {
|
||||
self.handle_scrape_request(request).await?;
|
||||
}
|
||||
Err(err) => {
|
||||
::log::debug!("Couldn't parse in_message: {:?}", err);
|
||||
::log::debug!("Couldn't parse in_message: {:#}", err);
|
||||
|
||||
self.send_error_response("Invalid request".into(), None, None)
|
||||
.await?;
|
||||
|
|
@ -261,9 +264,7 @@ impl<S: futures::AsyncRead + futures::AsyncWrite + Unpin> ConnectionReader<S> {
|
|||
}
|
||||
}
|
||||
|
||||
async fn handle_in_message(&mut self, in_message: InMessage) -> anyhow::Result<()> {
|
||||
match in_message {
|
||||
InMessage::AnnounceRequest(announce_request) => {
|
||||
async fn handle_announce_request(&mut self, request: AnnounceRequest) -> anyhow::Result<()> {
|
||||
#[cfg(feature = "metrics")]
|
||||
::metrics::increment_counter!(
|
||||
"aquatic_requests_total",
|
||||
|
|
@ -272,20 +273,19 @@ impl<S: futures::AsyncRead + futures::AsyncWrite + Unpin> ConnectionReader<S> {
|
|||
"worker_index" => WORKER_INDEX.with(|index| index.get()).to_string(),
|
||||
);
|
||||
|
||||
let info_hash = announce_request.info_hash;
|
||||
let info_hash = request.info_hash;
|
||||
|
||||
if self
|
||||
.access_list_cache
|
||||
.load()
|
||||
.allows(self.config.access_list.mode, &info_hash.0)
|
||||
{
|
||||
let mut announced_info_hashes =
|
||||
self.clean_up_data.announced_info_hashes.borrow_mut();
|
||||
let mut announced_info_hashes = self.clean_up_data.announced_info_hashes.borrow_mut();
|
||||
|
||||
// Store peer id / check if stored peer id matches
|
||||
match announced_info_hashes.entry(announce_request.info_hash) {
|
||||
match announced_info_hashes.entry(request.info_hash) {
|
||||
Entry::Occupied(entry) => {
|
||||
if *entry.get() != announce_request.peer_id {
|
||||
if *entry.get() != request.peer_id {
|
||||
// Drop Rc borrow before awaiting
|
||||
drop(announced_info_hashes);
|
||||
|
||||
|
|
@ -302,7 +302,7 @@ impl<S: futures::AsyncRead + futures::AsyncWrite + Unpin> ConnectionReader<S> {
|
|||
}
|
||||
}
|
||||
Entry::Vacant(entry) => {
|
||||
entry.insert(announce_request.peer_id);
|
||||
entry.insert(request.peer_id);
|
||||
|
||||
// Set peer client info if not set
|
||||
#[cfg(feature = "metrics")]
|
||||
|
|
@ -310,7 +310,7 @@ impl<S: futures::AsyncRead + futures::AsyncWrite + Unpin> ConnectionReader<S> {
|
|||
&& self.config.metrics.peer_clients
|
||||
&& self.clean_up_data.opt_peer_client.borrow().is_none()
|
||||
{
|
||||
let peer_id = aquatic_peer_id::PeerId(announce_request.peer_id.0);
|
||||
let peer_id = aquatic_peer_id::PeerId(request.peer_id.0);
|
||||
let client = peer_id.client();
|
||||
let prefix = peer_id.first_8_bytes_hex().to_string();
|
||||
|
||||
|
|
@ -328,23 +328,21 @@ impl<S: futures::AsyncRead + futures::AsyncWrite + Unpin> ConnectionReader<S> {
|
|||
);
|
||||
}
|
||||
|
||||
*self.clean_up_data.opt_peer_client.borrow_mut() =
|
||||
Some((client, prefix));
|
||||
*self.clean_up_data.opt_peer_client.borrow_mut() = Some((client, prefix));
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(AnnounceEvent::Stopped) = announce_request.event {
|
||||
announced_info_hashes.remove(&announce_request.info_hash);
|
||||
if let Some(AnnounceEvent::Stopped) = request.event {
|
||||
announced_info_hashes.remove(&request.info_hash);
|
||||
}
|
||||
|
||||
// Drop Rc borrow before awaiting
|
||||
drop(announced_info_hashes);
|
||||
|
||||
let in_message = InMessage::AnnounceRequest(announce_request);
|
||||
let in_message = InMessage::AnnounceRequest(request);
|
||||
|
||||
let consumer_index =
|
||||
calculate_in_message_consumer_index(&self.config, info_hash);
|
||||
let consumer_index = calculate_in_message_consumer_index(&self.config, info_hash);
|
||||
|
||||
// Only fails when receiver is closed
|
||||
self.in_message_senders
|
||||
|
|
@ -362,8 +360,11 @@ impl<S: futures::AsyncRead + futures::AsyncWrite + Unpin> ConnectionReader<S> {
|
|||
)
|
||||
.await?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
InMessage::ScrapeRequest(ScrapeRequest { info_hashes, .. }) => {
|
||||
|
||||
async fn handle_scrape_request(&mut self, request: ScrapeRequest) -> anyhow::Result<()> {
|
||||
#[cfg(feature = "metrics")]
|
||||
::metrics::increment_counter!(
|
||||
"aquatic_requests_total",
|
||||
|
|
@ -372,7 +373,7 @@ impl<S: futures::AsyncRead + futures::AsyncWrite + Unpin> ConnectionReader<S> {
|
|||
"worker_index" => WORKER_INDEX.with(|index| index.get()).to_string(),
|
||||
);
|
||||
|
||||
let info_hashes = if let Some(info_hashes) = info_hashes {
|
||||
let info_hashes = if let Some(info_hashes) = request.info_hashes {
|
||||
info_hashes
|
||||
} else {
|
||||
// If request.info_hashes is empty, don't return scrape for all
|
||||
|
|
@ -425,8 +426,6 @@ impl<S: futures::AsyncRead + futures::AsyncWrite + Unpin> ConnectionReader<S> {
|
|||
.await
|
||||
.unwrap();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
|
@ -485,31 +484,28 @@ impl<S: futures::AsyncRead + futures::AsyncWrite + Unpin> ConnectionWriter<S> {
|
|||
.pending_scrape_id
|
||||
.expect("meta.pending_scrape_id not set");
|
||||
|
||||
let finished = if let Some(pending) = Slab::get_mut(
|
||||
&mut RefCell::borrow_mut(&self.pending_scrape_slab),
|
||||
pending_scrape_id.0 as usize,
|
||||
) {
|
||||
pending.stats.extend(out_message.files);
|
||||
pending.pending_worker_out_messages -= 1;
|
||||
let mut pending_responses = self.pending_scrape_slab.borrow_mut();
|
||||
|
||||
pending.pending_worker_out_messages == 0
|
||||
} else {
|
||||
return Err(anyhow::anyhow!("pending scrape not found in slab"));
|
||||
};
|
||||
let pending_response = pending_responses
|
||||
.get_mut(pending_scrape_id.0 as usize)
|
||||
.ok_or(anyhow::anyhow!("pending scrape not found in slab"))?;
|
||||
|
||||
if finished {
|
||||
let out_message = {
|
||||
let mut slab = RefCell::borrow_mut(&self.pending_scrape_slab);
|
||||
pending_response.stats.extend(out_message.files);
|
||||
pending_response.pending_worker_out_messages -= 1;
|
||||
|
||||
let pending = slab.remove(pending_scrape_id.0 as usize);
|
||||
if pending_response.pending_worker_out_messages == 0 {
|
||||
let pending_response =
|
||||
pending_responses.remove(pending_scrape_id.0 as usize);
|
||||
|
||||
slab.shrink_to_fit();
|
||||
pending_responses.shrink_to_fit();
|
||||
|
||||
OutMessage::ScrapeResponse(ScrapeResponse {
|
||||
let out_message = OutMessage::ScrapeResponse(ScrapeResponse {
|
||||
action: ScrapeAction::Scrape,
|
||||
files: pending.stats,
|
||||
})
|
||||
};
|
||||
files: pending_response.stats,
|
||||
});
|
||||
|
||||
// Drop Rc borrow before awaiting
|
||||
drop(pending_responses);
|
||||
|
||||
self.send_out_message(&out_message).await?;
|
||||
}
|
||||
|
|
@ -522,18 +518,16 @@ impl<S: futures::AsyncRead + futures::AsyncWrite + Unpin> ConnectionWriter<S> {
|
|||
}
|
||||
|
||||
async fn send_out_message(&mut self, out_message: &OutMessage) -> anyhow::Result<()> {
|
||||
let result = timeout(Duration::from_secs(10), async {
|
||||
let result =
|
||||
futures::SinkExt::send(&mut self.ws_out, out_message.to_ws_message()).await;
|
||||
|
||||
Ok(result)
|
||||
timeout(Duration::from_secs(10), async {
|
||||
Ok(futures::SinkExt::send(&mut self.ws_out, out_message.to_ws_message()).await)
|
||||
})
|
||||
.await;
|
||||
.await
|
||||
.map_err(|err| {
|
||||
anyhow::anyhow!("send_out_message: sending to peer took too long: {:#}", err)
|
||||
})?
|
||||
.with_context(|| "send_out_message")?;
|
||||
|
||||
match result {
|
||||
Ok(Ok(())) => {
|
||||
if let OutMessage::AnnounceResponse(_) | OutMessage::ScrapeResponse(_) = out_message
|
||||
{
|
||||
if let OutMessage::AnnounceResponse(_) | OutMessage::ScrapeResponse(_) = out_message {
|
||||
*self.connection_valid_until.borrow_mut() = ValidUntil::new(
|
||||
self.server_start_instant,
|
||||
self.config.cleaning.max_connection_idle,
|
||||
|
|
@ -582,13 +576,6 @@ impl<S: futures::AsyncRead + futures::AsyncWrite + Unpin> ConnectionWriter<S> {
|
|||
|
||||
Ok(())
|
||||
}
|
||||
Ok(Err(err)) => Err(err.into()),
|
||||
Err(err) => Err(anyhow::anyhow!(
|
||||
"send_out_message: sending to peer took too long: {}",
|
||||
err
|
||||
)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Data stored with connection needed for cleanup after it closes
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue