From fe5ccf66460e1a3289cd17ac329dd3cae6324e1b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Joakim=20Frosteg=C3=A5rd?= Date: Thu, 9 Nov 2023 18:06:21 +0100 Subject: [PATCH] ws: refactor, bug fixes, improvements (#155) - split swarm worker into two modules - split socket worker into two modules - keep track of which offers peers have sent and only allow matching answers - always clean up after closing connection - use channel for telling connections to close - move some logic into new ConnectionRunner struct - use slotmap for connection reference storage - fix double counting of error responses - actually close connections that take too long to send responses to - remove announced_info_hashes entry on AnnounceEvent::Stopped --- CHANGELOG.md | 13 + Cargo.lock | 11 + TODO.md | 3 - crates/common/src/lib.rs | 10 +- crates/http/src/workers/swarm.rs | 6 +- crates/udp/src/workers/swarm/storage.rs | 8 +- crates/ws/Cargo.toml | 2 + crates/ws/src/common.rs | 5 +- crates/ws/src/config.rs | 7 +- crates/ws/src/lib.rs | 2 +- crates/ws/src/workers/socket.rs | 1010 -------------------- crates/ws/src/workers/socket/connection.rs | 667 +++++++++++++ crates/ws/src/workers/socket/mod.rs | 362 +++++++ crates/ws/src/workers/swarm.rs | 541 ----------- crates/ws/src/workers/swarm/mod.rs | 183 ++++ crates/ws/src/workers/swarm/storage.rs | 488 ++++++++++ crates/ws_load_test/src/network.rs | 1 + crates/ws_load_test/src/utils.rs | 32 +- crates/ws_protocol/src/common.rs | 2 +- 19 files changed, 1770 insertions(+), 1583 deletions(-) delete mode 100644 crates/ws/src/workers/socket.rs create mode 100644 crates/ws/src/workers/socket/connection.rs create mode 100644 crates/ws/src/workers/socket/mod.rs delete mode 100644 crates/ws/src/workers/swarm.rs create mode 100644 crates/ws/src/workers/swarm/mod.rs create mode 100644 crates/ws/src/workers/swarm/storage.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index b38bbdb..ddefb9d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -26,11 +26,24 @@ * Add support for reporting peer client information * Reload TLS certificate (and key) on SIGUSR1 +* Keep track of which offers peers have sent and only allow matching answers #### Changed +* A response is no longer generated when peers announce with AnnounceEvent::Stopped * Compiling with SIMD extensions enabled is no longer required, due to the addition of runtime detection to simd-json +* Only consider announce and scrape responses as signs of connection still + being alive. Previously, all messages sent to peer were considered. +* Decrease default max_peer_age and max_connection_idle config values + +#### Fixed + +* Fix bug where clean up after closing connections wasn't always done +* Fix double counting of error responses +* Actually close connections that are too slow to send responses to +* If peers announce with AnnounceEvent::Stopped, allow them to later announce on + same torrent with different peer_id ## 0.8.0 - 2023-03-17 diff --git a/Cargo.lock b/Cargo.lock index 7a09a85..2b2a24c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -321,6 +321,7 @@ dependencies = [ "glommio", "hashbrown 0.14.1", "httparse", + "indexmap 2.0.2", "log", "metrics", "metrics-exporter-prometheus", @@ -335,6 +336,7 @@ dependencies = [ "serde", "signal-hook", "slab", + "slotmap", "socket2 0.5.4", "tungstenite", ] @@ -2491,6 +2493,15 @@ dependencies = [ "autocfg", ] +[[package]] +name = "slotmap" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e1e08e261d0e8f5c43123b7adf3e4ca1690d655377ac93a03b2c9d3e98de1342" +dependencies = [ + "version_check", +] + [[package]] name = "smallvec" version = "1.11.1" diff --git a/TODO.md b/TODO.md index 3a23817..da43bc6 100644 --- a/TODO.md +++ b/TODO.md @@ -3,9 +3,6 @@ ## High priority * aquatic_ws - * Store OfferId that a peer has sent out and only allow answers matching - them to be sent? HashMap<(OfferId, PeerId), ValidUntil> could work, but - not if peers reuse offer ids * Validate SDP data ## Medium priority diff --git a/crates/common/src/lib.rs b/crates/common/src/lib.rs index 9842735..44f25f5 100644 --- a/crates/common/src/lib.rs +++ b/crates/common/src/lib.rs @@ -156,7 +156,7 @@ pub fn extract_response_peers( ) -> Vec where K: Eq + ::std::hash::Hash, - F: Fn(&V) -> R, + F: Fn(&K, &V) -> R, { if peer_map.len() <= max_num_peers_to_take + 1 { // This branch: number of peers in map (minus sender peer) is less than @@ -165,7 +165,7 @@ where let mut peers = Vec::with_capacity(peer_map.len()); peers.extend(peer_map.iter().filter_map(|(k, v)| { - (*k != sender_peer_map_key).then_some(peer_conversion_function(v)) + (*k != sender_peer_map_key).then_some(peer_conversion_function(k, v)) })); // Handle the case when sender peer is not in peer list. Typically, @@ -204,12 +204,12 @@ where if let Some(slice) = peer_map.get_range(offset_half_one..end_half_one) { peers.extend(slice.iter().filter_map(|(k, v)| { - (*k != sender_peer_map_key).then_some(peer_conversion_function(v)) + (*k != sender_peer_map_key).then_some(peer_conversion_function(k, v)) })); } if let Some(slice) = peer_map.get_range(offset_half_two..end_half_two) { peers.extend(slice.iter().filter_map(|(k, v)| { - (*k != sender_peer_map_key).then_some(peer_conversion_function(v)) + (*k != sender_peer_map_key).then_some(peer_conversion_function(k, v)) })); } @@ -260,7 +260,7 @@ mod tests { &peer_map, max_num_peers_to_take, sender_peer_map_key, - |p| *p, + |_, p| *p, ); if num_peers_in_map > max_num_peers_to_take + 1 { diff --git a/crates/http/src/workers/swarm.rs b/crates/http/src/workers/swarm.rs index a6a5aeb..bed4f55 100644 --- a/crates/http/src/workers/swarm.rs +++ b/crates/http/src/workers/swarm.rs @@ -79,10 +79,10 @@ pub struct Peer { } impl Peer { - pub fn to_response_peer(&self) -> ResponsePeer { + pub fn to_response_peer(_: &PeerMapKey, peer: &Self) -> ResponsePeer { ResponsePeer { - ip_address: self.ip_address, - port: self.port, + ip_address: peer.ip_address, + port: peer.port, } } } diff --git a/crates/udp/src/workers/swarm/storage.rs b/crates/udp/src/workers/swarm/storage.rs index 777dba4..567cc80 100644 --- a/crates/udp/src/workers/swarm/storage.rs +++ b/crates/udp/src/workers/swarm/storage.rs @@ -30,10 +30,10 @@ struct Peer { } impl Peer { - fn to_response_peer(&self) -> ResponsePeer { + fn to_response_peer(_: &PeerId, peer: &Self) -> ResponsePeer { ResponsePeer { - ip_address: self.ip_address, - port: self.port, + ip_address: peer.ip_address, + port: peer.port, } } } @@ -352,7 +352,7 @@ mod tests { if i == 0 { opt_sender_key = Some(key); - opt_sender_peer = Some(peer.to_response_peer()); + opt_sender_peer = Some(Peer::to_response_peer(&key, &peer)); } peer_map.insert(key, peer); diff --git a/crates/ws/Cargo.toml b/crates/ws/Cargo.toml index a05ee15..2d7d8c8 100644 --- a/crates/ws/Cargo.toml +++ b/crates/ws/Cargo.toml @@ -38,6 +38,7 @@ futures-rustls = "0.24" glommio = "0.8" hashbrown = { version = "0.14", features = ["serde"] } httparse = "1" +indexmap = "2" log = "0.4" metrics = { version = "0.21", optional = true } metrics-util = { version = "0.15", optional = true } @@ -50,6 +51,7 @@ rustls-pemfile = "1" serde = { version = "1", features = ["derive"] } signal-hook = { version = "0.3" } slab = "0.4" +slotmap = "1" socket2 = { version = "0.5", features = ["all"] } tungstenite = "0.20" diff --git a/crates/ws/src/common.rs b/crates/ws/src/common.rs index d4e6d35..dae682f 100644 --- a/crates/ws/src/common.rs +++ b/crates/ws/src/common.rs @@ -34,8 +34,9 @@ pub struct PendingScrapeId(pub u8); #[derive(Copy, Clone, Debug)] pub struct ConsumerId(pub u8); -#[derive(Clone, Copy, Debug, PartialEq)] -pub struct ConnectionId(pub usize); +slotmap::new_key_type! { + pub struct ConnectionId; +} #[derive(Clone, Copy, Debug)] pub struct InMessageMeta { diff --git a/crates/ws/src/config.rs b/crates/ws/src/config.rs index 301d3fc..e2a9376 100644 --- a/crates/ws/src/config.rs +++ b/crates/ws/src/config.rs @@ -154,6 +154,8 @@ pub struct CleaningConfig { pub torrent_cleaning_interval: u64, /// Remove peers that have not announced for this long (seconds) pub max_peer_age: u32, + /// Require that offers are answered to withing this period (seconds) + pub max_offer_age: u32, // Clean connections this often (seconds) pub connection_cleaning_interval: u64, /// Close connections if no responses have been sent to them for this long (seconds) @@ -169,8 +171,9 @@ impl Default for CleaningConfig { fn default() -> Self { Self { torrent_cleaning_interval: 30, - max_peer_age: 1800, - max_connection_idle: 60 * 5, + max_peer_age: 180, + max_offer_age: 120, + max_connection_idle: 180, connection_cleaning_interval: 30, close_after_tls_update_grace_period: 60 * 60 * 60, } diff --git a/crates/ws/src/lib.rs b/crates/ws/src/lib.rs index 0a536cd..806267f 100644 --- a/crates/ws/src/lib.rs +++ b/crates/ws/src/lib.rs @@ -71,7 +71,7 @@ pub fn run(config: Config) -> ::anyhow::Result<()> { let request_mesh_builder = MeshBuilder::partial(num_peers, SHARED_IN_CHANNEL_SIZE); let response_mesh_builder = MeshBuilder::partial(num_peers, SHARED_IN_CHANNEL_SIZE * 16); - let control_mesh_builder = MeshBuilder::partial(num_peers, SHARED_IN_CHANNEL_SIZE); + let control_mesh_builder = MeshBuilder::partial(num_peers, SHARED_IN_CHANNEL_SIZE * 16); let (sentinel_watcher, sentinel) = PanicSentinelWatcher::create_with_sentinel(); let priv_dropper = PrivilegeDropper::new(config.privileges.clone(), config.socket_workers); diff --git a/crates/ws/src/workers/socket.rs b/crates/ws/src/workers/socket.rs deleted file mode 100644 index a53e9f0..0000000 --- a/crates/ws/src/workers/socket.rs +++ /dev/null @@ -1,1010 +0,0 @@ -use std::borrow::Cow; -use std::cell::RefCell; -use std::collections::BTreeMap; -use std::os::unix::prelude::{FromRawFd, IntoRawFd}; -use std::rc::Rc; -use std::sync::Arc; -use std::time::Duration; - -use anyhow::Context; -use aquatic_common::access_list::{create_access_list_cache, AccessListArcSwap, AccessListCache}; -use aquatic_common::privileges::PrivilegeDropper; -use aquatic_common::rustls_config::RustlsConfig; -use aquatic_common::{PanicSentinel, ServerStartInstant}; -use aquatic_peer_id::PeerClient; -use aquatic_ws_protocol::*; -use arc_swap::ArcSwap; -use async_tungstenite::WebSocketStream; -use futures::stream::{SplitSink, SplitStream}; -use futures::{AsyncWriteExt, StreamExt}; -use futures_lite::future::race; -use futures_rustls::TlsAcceptor; -use glommio::channels::channel_mesh::{MeshBuilder, Partial, Role, Senders}; -use glommio::channels::local_channel::{new_bounded, LocalReceiver, LocalSender}; -use glommio::channels::shared_channel::ConnectedReceiver; -use glommio::net::{TcpListener, TcpStream}; -use glommio::task::JoinHandle; -use glommio::timer::{sleep, timeout, TimerActionRepeat}; -use glommio::{enclose, prelude::*}; -use hashbrown::hash_map::Entry; -use hashbrown::HashMap; -use slab::Slab; - -use crate::config::Config; - -use crate::common::*; - -const LOCAL_CHANNEL_SIZE: usize = 16; - -#[cfg(feature = "metrics")] -thread_local! { static WORKER_INDEX: ::std::cell::Cell = Default::default() } - -struct PendingScrapeResponse { - pending_worker_out_messages: usize, - stats: HashMap, -} - -struct ConnectionReference { - task_handle: Option>, - /// Sender part of channel used to pass on outgoing messages from request - /// worker - out_message_sender: Rc>, - /// Updated after sending message to peer - valid_until: ValidUntil, - announced_info_hashes: HashMap, - ip_version: IpVersion, - opt_peer_client: Option<(PeerClient, String)>, - opt_tls_config: Option>, - valid_until_after_tls_update: Option, -} - -pub async fn run_socket_worker( - _sentinel: PanicSentinel, - config: Config, - state: State, - opt_tls_config: Option>>, - control_message_mesh_builder: MeshBuilder, - in_message_mesh_builder: MeshBuilder<(InMessageMeta, InMessage), Partial>, - out_message_mesh_builder: MeshBuilder<(OutMessageMeta, OutMessage), Partial>, - priv_dropper: PrivilegeDropper, - server_start_instant: ServerStartInstant, - worker_index: usize, -) { - #[cfg(feature = "metrics")] - WORKER_INDEX.with(|index| index.set(worker_index)); - - let config = Rc::new(config); - let access_list = state.access_list; - - let listener = create_tcp_listener(&config, priv_dropper).expect("create tcp listener"); - - ::log::info!("created tcp listener"); - - let (control_message_senders, _) = control_message_mesh_builder - .join(Role::Producer) - .await - .unwrap(); - let control_message_senders = Rc::new(control_message_senders); - - let (in_message_senders, _) = in_message_mesh_builder.join(Role::Producer).await.unwrap(); - let in_message_senders = Rc::new(in_message_senders); - - let tq_prioritized = executor().create_task_queue( - Shares::Static(100), - Latency::Matters(Duration::from_millis(1)), - "prioritized", - ); - let tq_regular = - executor().create_task_queue(Shares::Static(1), Latency::NotImportant, "regular"); - - let (_, mut out_message_receivers) = - out_message_mesh_builder.join(Role::Consumer).await.unwrap(); - let out_message_consumer_id = ConsumerId( - out_message_receivers - .consumer_id() - .unwrap() - .try_into() - .unwrap(), - ); - - ::log::info!("joined channels"); - - let connection_slab = Rc::new(RefCell::new(Slab::new())); - - // Periodically clean connections - TimerActionRepeat::repeat_into( - enclose!((config, connection_slab, opt_tls_config) move || { - clean_connections( - config.clone(), - connection_slab.clone(), - server_start_instant, - opt_tls_config.clone(), - ) - }), - tq_prioritized, - ) - .unwrap(); - - for (_, out_message_receiver) in out_message_receivers.streams() { - spawn_local_into( - receive_out_messages(out_message_receiver, connection_slab.clone()), - tq_regular, - ) - .unwrap() - .detach(); - } - - let mut incoming = listener.incoming(); - - while let Some(stream) = incoming.next().await { - match stream { - Ok(stream) => { - let ip_version = match stream.peer_addr() { - Ok(addr) => IpVersion::canonical_from_ip(addr.ip()), - Err(err) => { - ::log::info!("could not extract ip version (v4 or v6): {:#}", err); - - continue; - } - }; - - let (out_message_sender, out_message_receiver) = new_bounded(LOCAL_CHANNEL_SIZE); - let out_message_sender = Rc::new(out_message_sender); - - let key = RefCell::borrow_mut(&connection_slab).insert(ConnectionReference { - task_handle: None, - out_message_sender: out_message_sender.clone(), - valid_until: ValidUntil::new( - server_start_instant, - config.cleaning.max_connection_idle, - ), - announced_info_hashes: Default::default(), - ip_version, - opt_peer_client: None, - opt_tls_config: opt_tls_config.as_ref().map(|c| c.load_full()), - valid_until_after_tls_update: None, - }); - - ::log::trace!("accepting stream, assigning id {}", key); - - let task_handle = spawn_local_into(enclose!((config, access_list, control_message_senders, in_message_senders, connection_slab, opt_tls_config) async move { - #[cfg(feature = "metrics")] - ::metrics::increment_gauge!( - "aquatic_active_connections", - 1.0, - "ip_version" => ip_version_to_metrics_str(ip_version), - "worker_index" => worker_index.to_string(), - ); - - if let Err(err) = run_connection( - config.clone(), - access_list, - in_message_senders, - tq_prioritized, - tq_regular, - connection_slab.clone(), - out_message_sender, - out_message_receiver, - server_start_instant, - out_message_consumer_id, - ConnectionId(key), - opt_tls_config, - ip_version, - stream, - ).await { - ::log::debug!("connection error: {:#}", err); - } - - // Clean up after closed connection - - #[cfg(feature = "metrics")] - ::metrics::decrement_gauge!( - "aquatic_active_connections", - 1.0, - "ip_version" => ip_version_to_metrics_str(ip_version), - "worker_index" => worker_index.to_string(), - ); - - // Remove reference in separate statement to avoid - // multiple RefCell borrows - let opt_reference = connection_slab.borrow_mut().try_remove(key); - - // Tell swarm workers to remove peer - if let Some(reference) = opt_reference { - for (info_hash, peer_id) in reference.announced_info_hashes { - let message = SwarmControlMessage::ConnectionClosed { - info_hash, - peer_id, - ip_version: reference.ip_version, - }; - - let consumer_index = - calculate_in_message_consumer_index(&config, info_hash); - - // Only fails when receiver is closed - control_message_senders - .send_to( - consumer_index, - message - ) - .await - .unwrap(); - } - - #[cfg(feature = "prometheus")] - if let Some((peer_client, prefix)) = reference.opt_peer_client { - ::metrics::decrement_gauge!( - "aquatic_peer_clients", - 1.0, - "client" => peer_client.to_string(), - ); - - if config.metrics.peer_id_prefixes { - ::metrics::decrement_gauge!( - "aquatic_peer_id_prefixes", - 1.0, - "prefix_hex" => prefix.to_string(), - ); - } - } - } - }), tq_regular) - .unwrap() - .detach(); - - if let Some(reference) = connection_slab.borrow_mut().get_mut(key) { - reference.task_handle = Some(task_handle); - } - } - Err(err) => { - ::log::error!("accept connection: {:#}", err); - } - } - } -} - -async fn clean_connections( - config: Rc, - connection_slab: Rc>>, - server_start_instant: ServerStartInstant, - opt_tls_config: Option>>, -) -> Option { - let now = server_start_instant.seconds_elapsed(); - let opt_current_tls_config = opt_tls_config.map(|c| c.load_full()); - - connection_slab.borrow_mut().retain(|_, reference| { - if let Some(valid_until) = reference.valid_until_after_tls_update { - if !valid_until.valid(now) { - if let Some(handle) = reference.task_handle.as_ref() { - handle.cancel(); - } - - return false; - } - } else if let Some(false) = opt_current_tls_config - .as_ref() - .zip(reference.opt_tls_config.as_ref()) - .map(|(a, b)| Arc::ptr_eq(a, b)) - { - reference.valid_until_after_tls_update = Some(ValidUntil::new( - server_start_instant, - config.cleaning.close_after_tls_update_grace_period, - )); - } - - if reference.valid_until.valid(now) { - #[cfg(feature = "prometheus")] - if let Some((peer_client, prefix)) = &reference.opt_peer_client { - // As long as connection is still alive, increment peer client - // gauges by zero to prevent them from being removed due to - // idleness - - ::metrics::increment_gauge!( - "aquatic_peer_clients", - 0.0, - "client" => peer_client.to_string(), - ); - - if config.metrics.peer_id_prefixes { - ::metrics::increment_gauge!( - "aquatic_peer_id_prefixes", - 0.0, - "prefix_hex" => prefix.to_string(), - ); - } - } - - true - } else { - if let Some(ref handle) = reference.task_handle { - handle.cancel(); - } - - false - } - }); - - connection_slab.borrow_mut().shrink_to_fit(); - - #[cfg(feature = "metrics")] - { - // Increment gauges by zero to prevent them from being removed due to - // idleness - - let worker_index = WORKER_INDEX.with(|index| index.get()).to_string(); - - if config.network.address.is_ipv4() || !config.network.only_ipv6 { - ::metrics::increment_gauge!( - "aquatic_active_connections", - 0.0, - "ip_version" => "4", - "worker_index" => worker_index.clone(), - ); - } - if config.network.address.is_ipv6() { - ::metrics::increment_gauge!( - "aquatic_active_connections", - 0.0, - "ip_version" => "6", - "worker_index" => worker_index, - ); - } - } - - Some(Duration::from_secs( - config.cleaning.connection_cleaning_interval, - )) -} - -async fn receive_out_messages( - mut out_message_receiver: ConnectedReceiver<(OutMessageMeta, OutMessage)>, - connection_references: Rc>>, -) { - let connection_references = &connection_references; - - while let Some((meta, out_message)) = out_message_receiver.next().await { - if let Some(reference) = connection_references.borrow().get(meta.connection_id.0) { - ::log::trace!( - "local channel {} len: {}", - meta.connection_id.0, - reference.out_message_sender.len() - ); - - match reference.out_message_sender.try_send((meta, out_message)) { - Ok(()) => {} - Err(GlommioError::Closed(_)) => {} - Err(GlommioError::WouldBlock(_)) => {} - Err(err) => { - ::log::debug!( - "Couldn't send out_message from shared channel to local receiver: {:?}", - err - ); - } - } - } - } -} - -async fn run_connection( - config: Rc, - access_list: Arc, - in_message_senders: Rc>, - tq_prioritized: TaskQueueHandle, - tq_regular: TaskQueueHandle, - connection_slab: Rc>>, - out_message_sender: Rc>, - out_message_receiver: LocalReceiver<(OutMessageMeta, OutMessage)>, - server_start_instant: ServerStartInstant, - out_message_consumer_id: ConsumerId, - connection_id: ConnectionId, - opt_tls_config: Option>>, - ip_version: IpVersion, - mut stream: TcpStream, -) -> anyhow::Result<()> { - if let Some(tls_config) = opt_tls_config { - let tls_acceptor: TlsAcceptor = tls_config.load_full().into(); - - let stream = tls_acceptor.accept(stream).await?; - - run_stream_agnostic_connection( - config.clone(), - access_list, - in_message_senders, - tq_prioritized, - tq_regular, - connection_slab.clone(), - out_message_sender, - out_message_receiver, - server_start_instant, - out_message_consumer_id, - connection_id, - stream, - ip_version, - ) - .await - } else { - // Implementing this over TLS is too cumbersome, since the crate used - // for TLS streams doesn't support peek and tungstenite doesn't - // properly support sending a HTTP error response in accept_hdr - // callback. - if config.network.enable_http_health_checks { - let mut peek_buf = [0u8; 11]; - - stream - .peek(&mut peek_buf) - .await - .map_err(|err| anyhow::anyhow!("error peeking: {:#}", err))?; - - if &peek_buf == b"GET /health" { - stream - .write_all(b"HTTP/1.1 200 Ok\r\nContent-Length: 2\r\n\r\nOk") - .await - .map_err(|err| { - anyhow::anyhow!("error sending health check response: {:#}", err) - })?; - stream.flush().await.map_err(|err| { - anyhow::anyhow!("error flushing health check response: {:#}", err) - })?; - - return Err(anyhow::anyhow!( - "client requested health check, skipping websocket negotiation" - )); - } - } - - run_stream_agnostic_connection( - config.clone(), - access_list, - in_message_senders, - tq_prioritized, - tq_regular, - connection_slab.clone(), - out_message_sender, - out_message_receiver, - server_start_instant, - out_message_consumer_id, - connection_id, - stream, - ip_version, - ) - .await - } -} - -async fn run_stream_agnostic_connection< - S: futures::AsyncRead + futures::AsyncWrite + Unpin + 'static, ->( - config: Rc, - access_list: Arc, - in_message_senders: Rc>, - tq_prioritized: TaskQueueHandle, - tq_regular: TaskQueueHandle, - connection_slab: Rc>>, - out_message_sender: Rc>, - out_message_receiver: LocalReceiver<(OutMessageMeta, OutMessage)>, - server_start_instant: ServerStartInstant, - out_message_consumer_id: ConsumerId, - connection_id: ConnectionId, - stream: S, - ip_version: IpVersion, -) -> anyhow::Result<()> { - let ws_config = tungstenite::protocol::WebSocketConfig { - max_frame_size: Some(config.network.websocket_max_frame_size), - max_message_size: Some(config.network.websocket_max_message_size), - write_buffer_size: config.network.websocket_write_buffer_size, - max_write_buffer_size: config.network.websocket_write_buffer_size * 3, - ..Default::default() - }; - let stream = async_tungstenite::accept_async_with_config(stream, Some(ws_config)).await?; - - let (ws_out, ws_in) = futures::StreamExt::split(stream); - - let pending_scrape_slab = Rc::new(RefCell::new(Slab::new())); - let access_list_cache = create_access_list_cache(&access_list); - - let reader_handle = spawn_local_into( - enclose!((config, connection_slab, pending_scrape_slab) async move { - let mut reader = ConnectionReader { - config, - access_list_cache, - connection_slab, - in_message_senders, - out_message_sender, - pending_scrape_slab, - out_message_consumer_id, - ws_in, - ip_version, - connection_id, - }; - - let result = reader.run_in_message_loop().await; - - result - }), - tq_regular, - ) - .unwrap() - .detach(); - - let writer_handle = spawn_local_into( - async move { - let mut writer = ConnectionWriter { - config, - out_message_receiver, - connection_slab, - ws_out, - pending_scrape_slab, - connection_id, - server_start_instant, - ip_version, - }; - - let result = writer.run_out_message_loop().await; - - result - }, - tq_prioritized, - ) - .unwrap() - .detach(); - - race(reader_handle, writer_handle) - .await - .ok_or_else(|| anyhow::anyhow!("reader/writer task closed"))? -} - -struct ConnectionReader { - config: Rc, - access_list_cache: AccessListCache, - connection_slab: Rc>>, - in_message_senders: Rc>, - out_message_sender: Rc>, - pending_scrape_slab: Rc>>, - out_message_consumer_id: ConsumerId, - ws_in: SplitStream>, - ip_version: IpVersion, - connection_id: ConnectionId, -} - -impl ConnectionReader { - async fn run_in_message_loop(&mut self) -> anyhow::Result<()> { - loop { - while self.out_message_sender.is_full() { - sleep(Duration::from_millis(100)).await; - - yield_if_needed().await; - } - - let message = self - .ws_in - .next() - .await - .ok_or_else(|| anyhow::anyhow!("Stream ended"))??; - - match &message { - tungstenite::Message::Text(_) | tungstenite::Message::Binary(_) => { - match InMessage::from_ws_message(message) { - Ok(in_message) => { - self.handle_in_message(in_message).await?; - } - Err(err) => { - ::log::debug!("Couldn't parse in_message: {:?}", err); - - self.send_error_response("Invalid request".into(), None, None) - .await?; - } - } - } - tungstenite::Message::Ping(_) => { - ::log::trace!("Received ping message"); - // tungstenite sends a pong response by itself - } - tungstenite::Message::Pong(_) => { - ::log::trace!("Received pong message"); - } - tungstenite::Message::Close(_) => { - ::log::debug!("Client sent close frame"); - - break Ok(()); - } - tungstenite::Message::Frame(_) => { - ::log::warn!("Read raw websocket frame, this should not happen"); - } - } - - yield_if_needed().await; - } - } - - async fn handle_in_message(&mut self, in_message: InMessage) -> anyhow::Result<()> { - match in_message { - InMessage::AnnounceRequest(announce_request) => { - #[cfg(feature = "metrics")] - ::metrics::increment_counter!( - "aquatic_requests_total", - "type" => "announce", - "ip_version" => ip_version_to_metrics_str(self.ip_version), - "worker_index" => WORKER_INDEX.with(|index| index.get()).to_string(), - ); - - let info_hash = announce_request.info_hash; - - if self - .access_list_cache - .load() - .allows(self.config.access_list.mode, &info_hash.0) - { - { - let mut connection_slab = self.connection_slab.borrow_mut(); - - let connection_reference = connection_slab - .get_mut(self.connection_id.0) - .ok_or_else(|| { - anyhow::anyhow!( - "connection reference {} not found in slab", - self.connection_id.0 - ) - })?; - - // Store peer id / check if stored peer id matches - match connection_reference - .announced_info_hashes - .entry(announce_request.info_hash) - { - Entry::Occupied(entry) => { - if *entry.get() != announce_request.peer_id { - // Drop Rc borrow before awaiting - drop(connection_slab); - - self.send_error_response( - "Only one peer id can be used per torrent".into(), - Some(ErrorResponseAction::Announce), - Some(info_hash), - ) - .await?; - - return Err(anyhow::anyhow!( - "Peer used more than one PeerId for a single torrent" - )); - } - } - Entry::Vacant(entry) => { - #[cfg(feature = "prometheus")] - if self.config.metrics.run_prometheus_endpoint - && self.config.metrics.peer_clients - && connection_reference.opt_peer_client.is_none() - { - let peer_id = - aquatic_peer_id::PeerId(announce_request.peer_id.0); - let client = peer_id.client(); - let prefix = peer_id.first_8_bytes_hex().to_string(); - - ::metrics::increment_gauge!( - "aquatic_peer_clients", - 1.0, - "client" => client.to_string(), - ); - - if self.config.metrics.peer_id_prefixes { - ::metrics::increment_gauge!( - "aquatic_peer_id_prefixes", - 1.0, - "prefix_hex" => prefix.to_string(), - ); - } - - connection_reference.opt_peer_client = Some((client, prefix)); - }; - - entry.insert(announce_request.peer_id); - } - } - } - - let in_message = InMessage::AnnounceRequest(announce_request); - - let consumer_index = - calculate_in_message_consumer_index(&self.config, info_hash); - - // Only fails when receiver is closed - self.in_message_senders - .send_to( - consumer_index, - (self.make_connection_meta(None), in_message), - ) - .await - .unwrap(); - } else { - self.send_error_response( - "Info hash not allowed".into(), - Some(ErrorResponseAction::Announce), - Some(info_hash), - ) - .await?; - } - } - InMessage::ScrapeRequest(ScrapeRequest { info_hashes, .. }) => { - #[cfg(feature = "metrics")] - ::metrics::increment_counter!( - "aquatic_requests_total", - "type" => "scrape", - "ip_version" => ip_version_to_metrics_str(self.ip_version), - "worker_index" => WORKER_INDEX.with(|index| index.get()).to_string(), - ); - - let info_hashes = if let Some(info_hashes) = info_hashes { - info_hashes - } else { - // If request.info_hashes is empty, don't return scrape for all - // torrents, even though reference server does it. It is too expensive. - self.send_error_response( - "Full scrapes are not allowed".into(), - Some(ErrorResponseAction::Scrape), - None, - ) - .await?; - - return Ok(()); - }; - - let mut info_hashes_by_worker: BTreeMap> = BTreeMap::new(); - - for info_hash in info_hashes.as_vec() { - let info_hashes = info_hashes_by_worker - .entry(calculate_in_message_consumer_index(&self.config, info_hash)) - .or_default(); - - info_hashes.push(info_hash); - } - - let pending_worker_out_messages = info_hashes_by_worker.len(); - - let pending_scrape_response = PendingScrapeResponse { - pending_worker_out_messages, - stats: Default::default(), - }; - - let pending_scrape_id: u8 = self - .pending_scrape_slab - .borrow_mut() - .insert(pending_scrape_response) - .try_into() - .with_context(|| "Reached 256 pending scrape responses")?; - - let meta = self.make_connection_meta(Some(PendingScrapeId(pending_scrape_id))); - - for (consumer_index, info_hashes) in info_hashes_by_worker { - let in_message = InMessage::ScrapeRequest(ScrapeRequest { - action: ScrapeAction::Scrape, - info_hashes: Some(ScrapeRequestInfoHashes::Multiple(info_hashes)), - }); - - // Only fails when receiver is closed - self.in_message_senders - .send_to(consumer_index, (meta, in_message)) - .await - .unwrap(); - } - } - } - - Ok(()) - } - - async fn send_error_response( - &self, - failure_reason: Cow<'static, str>, - action: Option, - info_hash: Option, - ) -> anyhow::Result<()> { - let out_message = OutMessage::ErrorResponse(ErrorResponse { - action, - failure_reason, - info_hash, - }); - - let result = self - .out_message_sender - .send((self.make_connection_meta(None).into(), out_message)) - .await - .map_err(|err| { - anyhow::anyhow!("ConnectionReader::send_error_response failed: {}", err) - }); - - #[cfg(feature = "metrics")] - ::metrics::increment_counter!( - "aquatic_responses_total", - "type" => "error", - "ip_version" => ip_version_to_metrics_str(self.ip_version), - "worker_index" => WORKER_INDEX.with(|index| index.get()).to_string(), - ); - - result - } - - fn make_connection_meta(&self, pending_scrape_id: Option) -> InMessageMeta { - InMessageMeta { - connection_id: self.connection_id, - out_message_consumer_id: self.out_message_consumer_id, - ip_version: self.ip_version, - pending_scrape_id, - } - } -} - -struct ConnectionWriter { - config: Rc, - out_message_receiver: LocalReceiver<(OutMessageMeta, OutMessage)>, - connection_slab: Rc>>, - ws_out: SplitSink, tungstenite::Message>, - pending_scrape_slab: Rc>>, - server_start_instant: ServerStartInstant, - connection_id: ConnectionId, - ip_version: IpVersion, -} - -impl ConnectionWriter { - async fn run_out_message_loop(&mut self) -> anyhow::Result<()> { - loop { - let (meta, out_message) = self.out_message_receiver.recv().await.ok_or_else(|| { - anyhow::anyhow!("ConnectionWriter couldn't receive message, sender is closed") - })?; - - match out_message { - OutMessage::ScrapeResponse(out_message) => { - let pending_scrape_id = meta - .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; - - pending.pending_worker_out_messages == 0 - } else { - return Err(anyhow::anyhow!("pending scrape not found in slab")); - }; - - if finished { - let out_message = { - let mut slab = RefCell::borrow_mut(&self.pending_scrape_slab); - - let pending = slab.remove(pending_scrape_id.0 as usize); - - slab.shrink_to_fit(); - - OutMessage::ScrapeResponse(ScrapeResponse { - action: ScrapeAction::Scrape, - files: pending.stats, - }) - }; - - self.send_out_message(&out_message).await?; - } - } - out_message => { - self.send_out_message(&out_message).await?; - } - }; - } - } - - 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) - }) - .await; - - match result { - Ok(Ok(())) => { - #[cfg(feature = "metrics")] - { - let out_message_type = match &out_message { - OutMessage::OfferOutMessage(_) => "offer", - OutMessage::AnswerOutMessage(_) => "offer_answer", - OutMessage::AnnounceResponse(_) => "announce", - OutMessage::ScrapeResponse(_) => "scrape", - OutMessage::ErrorResponse(_) => "error", - }; - - ::metrics::increment_counter!( - "aquatic_responses_total", - "type" => out_message_type, - "ip_version" => ip_version_to_metrics_str(self.ip_version), - "worker_index" => WORKER_INDEX.with(|index| index.get()).to_string(), - ); - } - - self.connection_slab - .borrow_mut() - .get_mut(self.connection_id.0) - .ok_or_else(|| { - anyhow::anyhow!( - "connection reference {} not found in slab", - self.connection_id.0 - ) - })? - .valid_until = ValidUntil::new( - self.server_start_instant, - self.config.cleaning.max_connection_idle, - ); - - Ok(()) - } - Ok(Err(err)) => Err(err.into()), - Err(err) => { - ::log::debug!("send_out_message: sending to peer took too long: {}", err); - - Ok(()) - } - } - } -} - -fn calculate_in_message_consumer_index(config: &Config, info_hash: InfoHash) -> usize { - (info_hash.0[0] as usize) % config.swarm_workers -} - -fn create_tcp_listener( - config: &Config, - priv_dropper: PrivilegeDropper, -) -> anyhow::Result { - let domain = if config.network.address.is_ipv4() { - socket2::Domain::IPV4 - } else { - socket2::Domain::IPV6 - }; - - ::log::info!("creating socket.."); - - let socket = socket2::Socket::new(domain, socket2::Type::STREAM, Some(socket2::Protocol::TCP)) - .with_context(|| "create socket")?; - - if config.network.only_ipv6 { - ::log::info!("setting socket to ipv6 only.."); - - socket - .set_only_v6(true) - .with_context(|| "socket: set only ipv6")?; - } - - ::log::info!("setting SO_REUSEPORT on socket.."); - - socket - .set_reuse_port(true) - .with_context(|| "socket: set reuse port")?; - - ::log::info!("binding socket.."); - - socket - .bind(&config.network.address.into()) - .with_context(|| format!("socket: bind to {}", config.network.address))?; - - ::log::info!("listening on socket.."); - - socket - .listen(config.network.tcp_backlog) - .with_context(|| format!("socket: listen {}", config.network.address))?; - - ::log::info!("running PrivilegeDropper::after_socket_creation.."); - - priv_dropper.after_socket_creation()?; - - ::log::info!("casting socket to glommio TcpListener.."); - - Ok(unsafe { TcpListener::from_raw_fd(socket.into_raw_fd()) }) -} - -#[cfg(feature = "metrics")] -fn ip_version_to_metrics_str(ip_version: IpVersion) -> &'static str { - match ip_version { - IpVersion::V4 => "4", - IpVersion::V6 => "6", - } -} diff --git a/crates/ws/src/workers/socket/connection.rs b/crates/ws/src/workers/socket/connection.rs new file mode 100644 index 0000000..1232808 --- /dev/null +++ b/crates/ws/src/workers/socket/connection.rs @@ -0,0 +1,667 @@ +use std::borrow::Cow; +use std::cell::RefCell; +use std::collections::BTreeMap; +use std::rc::Rc; +use std::sync::Arc; +use std::time::Duration; + +use anyhow::Context; +use aquatic_common::access_list::{create_access_list_cache, AccessListArcSwap, AccessListCache}; +use aquatic_common::rustls_config::RustlsConfig; +use aquatic_common::ServerStartInstant; +use aquatic_peer_id::PeerClient; +use aquatic_ws_protocol::*; +use arc_swap::ArcSwap; +use async_tungstenite::WebSocketStream; +use futures::stream::{SplitSink, SplitStream}; +use futures::{AsyncWriteExt, StreamExt}; +use futures_lite::future::race; +use futures_rustls::TlsAcceptor; +use glommio::channels::channel_mesh::Senders; +use glommio::channels::local_channel::{LocalReceiver, LocalSender}; +use glommio::net::TcpStream; +use glommio::timer::{sleep, timeout}; +use glommio::{enclose, prelude::*}; +use hashbrown::hash_map::Entry; +use hashbrown::HashMap; +use slab::Slab; + +use crate::common::*; +use crate::config::Config; +use crate::workers::socket::calculate_in_message_consumer_index; + +#[cfg(feature = "metrics")] +use crate::workers::socket::{ip_version_to_metrics_str, WORKER_INDEX}; + +pub struct ConnectionRunner { + pub config: Rc, + pub access_list: Arc, + pub in_message_senders: Rc>, + pub tq_prioritized: TaskQueueHandle, + pub tq_regular: TaskQueueHandle, + pub connection_valid_until: Rc>, + pub out_message_sender: Rc>, + pub out_message_receiver: LocalReceiver<(OutMessageMeta, OutMessage)>, + pub close_conn_receiver: LocalReceiver<()>, + pub server_start_instant: ServerStartInstant, + pub out_message_consumer_id: ConsumerId, + pub connection_id: ConnectionId, + pub opt_tls_config: Option>>, + pub ip_version: IpVersion, +} + +impl ConnectionRunner { + pub async fn run( + self, + control_message_senders: Rc>, + stream: TcpStream, + ) { + let clean_up_data = ConnectionCleanupData { + announced_info_hashes: Default::default(), + ip_version: self.ip_version, + opt_peer_client: Default::default(), + }; + + clean_up_data.before_open(); + + let config = self.config.clone(); + + if let Err(err) = self.run_inner(clean_up_data.clone(), stream).await { + ::log::debug!("connection error: {:#}", err); + } + + clean_up_data + .after_close(&config, control_message_senders) + .await; + } + + async fn run_inner( + self, + clean_up_data: ConnectionCleanupData, + mut stream: TcpStream, + ) -> anyhow::Result<()> { + if let Some(tls_config) = self.opt_tls_config.as_ref() { + let tls_acceptor: TlsAcceptor = tls_config.load_full().into(); + + let stream = tls_acceptor.accept(stream).await?; + + self.run_inner_stream_agnostic(clean_up_data, stream).await + } else { + // Implementing this over TLS is too cumbersome, since the crate used + // for TLS streams doesn't support peek and tungstenite doesn't + // properly support sending a HTTP error response in accept_hdr + // callback. + if self.config.network.enable_http_health_checks { + let mut peek_buf = [0u8; 11]; + + stream + .peek(&mut peek_buf) + .await + .map_err(|err| anyhow::anyhow!("error peeking: {:#}", err))?; + + if &peek_buf == b"GET /health" { + stream + .write_all(b"HTTP/1.1 200 Ok\r\nContent-Length: 2\r\n\r\nOk") + .await + .map_err(|err| { + anyhow::anyhow!("error sending health check response: {:#}", err) + })?; + stream.flush().await.map_err(|err| { + anyhow::anyhow!("error flushing health check response: {:#}", err) + })?; + + return Err(anyhow::anyhow!( + "client requested health check, skipping websocket negotiation" + )); + } + } + + self.run_inner_stream_agnostic(clean_up_data, stream).await + } + } + + async fn run_inner_stream_agnostic( + self, + clean_up_data: ConnectionCleanupData, + stream: S, + ) -> anyhow::Result<()> + where + S: futures::AsyncRead + futures::AsyncWrite + Unpin + 'static, + { + let ws_config = tungstenite::protocol::WebSocketConfig { + max_frame_size: Some(self.config.network.websocket_max_frame_size), + max_message_size: Some(self.config.network.websocket_max_message_size), + write_buffer_size: self.config.network.websocket_write_buffer_size, + max_write_buffer_size: self.config.network.websocket_write_buffer_size * 3, + ..Default::default() + }; + let stream = async_tungstenite::accept_async_with_config(stream, Some(ws_config)).await?; + + let (ws_out, ws_in) = futures::StreamExt::split(stream); + + let pending_scrape_slab = Rc::new(RefCell::new(Slab::new())); + let access_list_cache = create_access_list_cache(&self.access_list); + + let config = self.config.clone(); + + let reader_handle = spawn_local_into( + enclose!((pending_scrape_slab, clean_up_data) async move { + let mut reader = ConnectionReader { + config: self.config.clone(), + access_list_cache, + in_message_senders: self.in_message_senders, + out_message_sender: self.out_message_sender, + pending_scrape_slab, + out_message_consumer_id: self.out_message_consumer_id, + ws_in, + ip_version: self.ip_version, + connection_id: self.connection_id, + clean_up_data: clean_up_data.clone(), + }; + + reader.run_in_message_loop().await + }), + self.tq_regular, + ) + .unwrap(); + + let writer_handle = spawn_local_into( + async move { + let mut writer = ConnectionWriter { + config, + out_message_receiver: self.out_message_receiver, + connection_valid_until: self.connection_valid_until, + ws_out, + pending_scrape_slab, + server_start_instant: self.server_start_instant, + ip_version: self.ip_version, + clean_up_data, + }; + + writer.run_out_message_loop().await + }, + self.tq_prioritized, + ) + .unwrap(); + + let close_conn_future = spawn_local_into( + async move { + self.close_conn_receiver.recv().await; + + Ok(()) + }, + self.tq_prioritized, + ) + .unwrap(); + + race(close_conn_future, race(reader_handle, writer_handle)).await + } +} + +struct ConnectionReader { + config: Rc, + access_list_cache: AccessListCache, + in_message_senders: Rc>, + out_message_sender: Rc>, + pending_scrape_slab: Rc>>, + out_message_consumer_id: ConsumerId, + ws_in: SplitStream>, + ip_version: IpVersion, + connection_id: ConnectionId, + clean_up_data: ConnectionCleanupData, +} + +impl ConnectionReader { + async fn run_in_message_loop(&mut self) -> anyhow::Result<()> { + loop { + while self.out_message_sender.is_full() { + sleep(Duration::from_millis(100)).await; + + yield_if_needed().await; + } + + let message = self + .ws_in + .next() + .await + .ok_or_else(|| anyhow::anyhow!("Stream ended"))??; + + match &message { + tungstenite::Message::Text(_) | tungstenite::Message::Binary(_) => { + match InMessage::from_ws_message(message) { + Ok(in_message) => { + self.handle_in_message(in_message).await?; + } + Err(err) => { + ::log::debug!("Couldn't parse in_message: {:?}", err); + + self.send_error_response("Invalid request".into(), None, None) + .await?; + } + } + } + tungstenite::Message::Ping(_) => { + ::log::trace!("Received ping message"); + // tungstenite sends a pong response by itself + } + tungstenite::Message::Pong(_) => { + ::log::trace!("Received pong message"); + } + tungstenite::Message::Close(_) => { + ::log::debug!("Client sent close frame"); + + break Ok(()); + } + tungstenite::Message::Frame(_) => { + ::log::warn!("Read raw websocket frame, this should not happen"); + } + } + + yield_if_needed().await; + } + } + + async fn handle_in_message(&mut self, in_message: InMessage) -> anyhow::Result<()> { + match in_message { + InMessage::AnnounceRequest(announce_request) => { + #[cfg(feature = "metrics")] + ::metrics::increment_counter!( + "aquatic_requests_total", + "type" => "announce", + "ip_version" => ip_version_to_metrics_str(self.ip_version), + "worker_index" => WORKER_INDEX.with(|index| index.get()).to_string(), + ); + + let info_hash = announce_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(); + + // Store peer id / check if stored peer id matches + match announced_info_hashes.entry(announce_request.info_hash) { + Entry::Occupied(entry) => { + if *entry.get() != announce_request.peer_id { + // Drop Rc borrow before awaiting + drop(announced_info_hashes); + + self.send_error_response( + "Only one peer id can be used per torrent".into(), + Some(ErrorResponseAction::Announce), + Some(info_hash), + ) + .await?; + + return Err(anyhow::anyhow!( + "Peer used more than one PeerId for a single torrent" + )); + } + } + Entry::Vacant(entry) => { + entry.insert(announce_request.peer_id); + + // Set peer client info if not set + #[cfg(feature = "metrics")] + if self.config.metrics.run_prometheus_endpoint + && 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 client = peer_id.client(); + let prefix = peer_id.first_8_bytes_hex().to_string(); + + ::metrics::increment_gauge!( + "aquatic_peer_clients", + 1.0, + "client" => client.to_string(), + ); + + if self.config.metrics.peer_id_prefixes { + ::metrics::increment_gauge!( + "aquatic_peer_id_prefixes", + 1.0, + "prefix_hex" => prefix.to_string(), + ); + } + + *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); + } + + // Drop Rc borrow before awaiting + drop(announced_info_hashes); + + let in_message = InMessage::AnnounceRequest(announce_request); + + let consumer_index = + calculate_in_message_consumer_index(&self.config, info_hash); + + // Only fails when receiver is closed + self.in_message_senders + .send_to( + consumer_index, + (self.make_connection_meta(None), in_message), + ) + .await + .unwrap(); + } else { + self.send_error_response( + "Info hash not allowed".into(), + Some(ErrorResponseAction::Announce), + Some(info_hash), + ) + .await?; + } + } + InMessage::ScrapeRequest(ScrapeRequest { info_hashes, .. }) => { + #[cfg(feature = "metrics")] + ::metrics::increment_counter!( + "aquatic_requests_total", + "type" => "scrape", + "ip_version" => ip_version_to_metrics_str(self.ip_version), + "worker_index" => WORKER_INDEX.with(|index| index.get()).to_string(), + ); + + let info_hashes = if let Some(info_hashes) = info_hashes { + info_hashes + } else { + // If request.info_hashes is empty, don't return scrape for all + // torrents, even though reference server does it. It is too expensive. + self.send_error_response( + "Full scrapes are not allowed".into(), + Some(ErrorResponseAction::Scrape), + None, + ) + .await?; + + return Ok(()); + }; + + let mut info_hashes_by_worker: BTreeMap> = BTreeMap::new(); + + for info_hash in info_hashes.as_vec() { + let info_hashes = info_hashes_by_worker + .entry(calculate_in_message_consumer_index(&self.config, info_hash)) + .or_default(); + + info_hashes.push(info_hash); + } + + let pending_worker_out_messages = info_hashes_by_worker.len(); + + let pending_scrape_response = PendingScrapeResponse { + pending_worker_out_messages, + stats: Default::default(), + }; + + let pending_scrape_id: u8 = self + .pending_scrape_slab + .borrow_mut() + .insert(pending_scrape_response) + .try_into() + .with_context(|| "Reached 256 pending scrape responses")?; + + let meta = self.make_connection_meta(Some(PendingScrapeId(pending_scrape_id))); + + for (consumer_index, info_hashes) in info_hashes_by_worker { + let in_message = InMessage::ScrapeRequest(ScrapeRequest { + action: ScrapeAction::Scrape, + info_hashes: Some(ScrapeRequestInfoHashes::Multiple(info_hashes)), + }); + + // Only fails when receiver is closed + self.in_message_senders + .send_to(consumer_index, (meta, in_message)) + .await + .unwrap(); + } + } + } + + Ok(()) + } + + async fn send_error_response( + &self, + failure_reason: Cow<'static, str>, + action: Option, + info_hash: Option, + ) -> anyhow::Result<()> { + let out_message = OutMessage::ErrorResponse(ErrorResponse { + action, + failure_reason, + info_hash, + }); + + self.out_message_sender + .send((self.make_connection_meta(None).into(), out_message)) + .await + .map_err(|err| { + anyhow::anyhow!("ConnectionReader::send_error_response failed: {:#}", err) + }) + } + + fn make_connection_meta(&self, pending_scrape_id: Option) -> InMessageMeta { + InMessageMeta { + connection_id: self.connection_id, + out_message_consumer_id: self.out_message_consumer_id, + ip_version: self.ip_version, + pending_scrape_id, + } + } +} + +struct ConnectionWriter { + config: Rc, + out_message_receiver: LocalReceiver<(OutMessageMeta, OutMessage)>, + connection_valid_until: Rc>, + ws_out: SplitSink, tungstenite::Message>, + pending_scrape_slab: Rc>>, + server_start_instant: ServerStartInstant, + ip_version: IpVersion, + clean_up_data: ConnectionCleanupData, +} + +impl ConnectionWriter { + async fn run_out_message_loop(&mut self) -> anyhow::Result<()> { + loop { + let (meta, out_message) = self.out_message_receiver.recv().await.ok_or_else(|| { + anyhow::anyhow!("ConnectionWriter couldn't receive message, sender is closed") + })?; + + match out_message { + OutMessage::ScrapeResponse(out_message) => { + let pending_scrape_id = meta + .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; + + pending.pending_worker_out_messages == 0 + } else { + return Err(anyhow::anyhow!("pending scrape not found in slab")); + }; + + if finished { + let out_message = { + let mut slab = RefCell::borrow_mut(&self.pending_scrape_slab); + + let pending = slab.remove(pending_scrape_id.0 as usize); + + slab.shrink_to_fit(); + + OutMessage::ScrapeResponse(ScrapeResponse { + action: ScrapeAction::Scrape, + files: pending.stats, + }) + }; + + self.send_out_message(&out_message).await?; + } + } + out_message => { + self.send_out_message(&out_message).await?; + } + }; + } + } + + 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) + }) + .await; + + match result { + Ok(Ok(())) => { + 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, + ); + } + + #[cfg(feature = "metrics")] + { + let out_message_type = match &out_message { + OutMessage::OfferOutMessage(_) => "offer", + OutMessage::AnswerOutMessage(_) => "offer_answer", + OutMessage::AnnounceResponse(_) => "announce", + OutMessage::ScrapeResponse(_) => "scrape", + OutMessage::ErrorResponse(_) => "error", + }; + + ::metrics::increment_counter!( + "aquatic_responses_total", + "type" => out_message_type, + "ip_version" => ip_version_to_metrics_str(self.ip_version), + "worker_index" => WORKER_INDEX.with(|index| index.get()).to_string(), + ); + + if let Some((peer_client, prefix)) = + self.clean_up_data.opt_peer_client.borrow().as_ref() + { + // As long as connection is still alive, increment peer client + // gauges by zero to prevent them from being removed due to + // idleness + + ::metrics::increment_gauge!( + "aquatic_peer_clients", + 0.0, + "client" => peer_client.to_string(), + ); + + if self.config.metrics.peer_id_prefixes { + ::metrics::increment_gauge!( + "aquatic_peer_id_prefixes", + 0.0, + "prefix_hex" => prefix.to_string(), + ); + } + } + } + + 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 +#[derive(Clone)] +struct ConnectionCleanupData { + announced_info_hashes: Rc>>, + ip_version: IpVersion, + opt_peer_client: Rc>>, +} + +impl ConnectionCleanupData { + fn before_open(&self) { + #[cfg(feature = "metrics")] + ::metrics::increment_gauge!( + "aquatic_active_connections", + 1.0, + "ip_version" => ip_version_to_metrics_str(self.ip_version), + "worker_index" => WORKER_INDEX.get().to_string(), + ); + } + async fn after_close( + &self, + config: &Config, + control_message_senders: Rc>, + ) { + // Use RefCell::take to avoid issues with Rc borrow across await + let announced_info_hashes = self.announced_info_hashes.take(); + + // Tell swarm workers to remove peer + for (info_hash, peer_id) in announced_info_hashes.into_iter() { + let message = SwarmControlMessage::ConnectionClosed { + info_hash, + peer_id, + ip_version: self.ip_version, + }; + + let consumer_index = calculate_in_message_consumer_index(&config, info_hash); + + control_message_senders + .send_to(consumer_index, message) + .await + .expect("control message receiver open"); + } + + #[cfg(feature = "metrics")] + { + ::metrics::decrement_gauge!( + "aquatic_active_connections", + 1.0, + "ip_version" => ip_version_to_metrics_str(self.ip_version), + "worker_index" => WORKER_INDEX.get().to_string(), + ); + + if let Some((peer_client, prefix)) = self.opt_peer_client.borrow().as_ref() { + ::metrics::decrement_gauge!( + "aquatic_peer_clients", + 1.0, + "client" => peer_client.to_string(), + ); + + if config.metrics.peer_id_prefixes { + ::metrics::decrement_gauge!( + "aquatic_peer_id_prefixes", + 1.0, + "prefix_hex" => prefix.to_string(), + ); + } + } + } + } +} + +struct PendingScrapeResponse { + pending_worker_out_messages: usize, + stats: HashMap, +} diff --git a/crates/ws/src/workers/socket/mod.rs b/crates/ws/src/workers/socket/mod.rs new file mode 100644 index 0000000..7bc6c72 --- /dev/null +++ b/crates/ws/src/workers/socket/mod.rs @@ -0,0 +1,362 @@ +use std::cell::RefCell; +use std::os::unix::prelude::{FromRawFd, IntoRawFd}; +use std::rc::Rc; +use std::sync::Arc; +use std::time::Duration; + +use anyhow::Context; +use aquatic_common::privileges::PrivilegeDropper; +use aquatic_common::rustls_config::RustlsConfig; +use aquatic_common::{PanicSentinel, ServerStartInstant}; +use aquatic_ws_protocol::*; +use arc_swap::ArcSwap; +use futures::StreamExt; +use glommio::channels::channel_mesh::{MeshBuilder, Partial, Role}; +use glommio::channels::local_channel::{new_bounded, LocalSender}; +use glommio::channels::shared_channel::ConnectedReceiver; +use glommio::net::TcpListener; +use glommio::timer::TimerActionRepeat; +use glommio::{enclose, prelude::*}; +use slotmap::HopSlotMap; + +use crate::config::Config; + +use crate::common::*; +use crate::workers::socket::connection::ConnectionRunner; + +mod connection; + +type ConnectionHandles = HopSlotMap; + +const LOCAL_CHANNEL_SIZE: usize = 16; + +#[cfg(feature = "metrics")] +thread_local! { static WORKER_INDEX: ::std::cell::Cell = Default::default() } + +/// Used to interact with the connection tasks +struct ConnectionHandle { + close_conn_sender: LocalSender<()>, + /// Sender part of channel used to pass on outgoing messages from request + /// worker + out_message_sender: Rc>, + /// Updated after sending message to peer + valid_until: Rc>, + /// The TLS config used for this connection + opt_tls_config: Option>, + valid_until_after_tls_update: Option, +} + +pub async fn run_socket_worker( + _sentinel: PanicSentinel, + config: Config, + state: State, + opt_tls_config: Option>>, + control_message_mesh_builder: MeshBuilder, + in_message_mesh_builder: MeshBuilder<(InMessageMeta, InMessage), Partial>, + out_message_mesh_builder: MeshBuilder<(OutMessageMeta, OutMessage), Partial>, + priv_dropper: PrivilegeDropper, + server_start_instant: ServerStartInstant, + worker_index: usize, +) { + #[cfg(feature = "metrics")] + WORKER_INDEX.with(|index| index.set(worker_index)); + + let config = Rc::new(config); + let access_list = state.access_list; + + let listener = create_tcp_listener(&config, priv_dropper).expect("create tcp listener"); + + ::log::info!("created tcp listener"); + + let (control_message_senders, _) = control_message_mesh_builder + .join(Role::Producer) + .await + .unwrap(); + let control_message_senders = Rc::new(control_message_senders); + + let (in_message_senders, _) = in_message_mesh_builder.join(Role::Producer).await.unwrap(); + let in_message_senders = Rc::new(in_message_senders); + + let tq_prioritized = executor().create_task_queue( + Shares::Static(100), + Latency::Matters(Duration::from_millis(1)), + "prioritized", + ); + let tq_regular = + executor().create_task_queue(Shares::Static(1), Latency::NotImportant, "regular"); + + let (_, mut out_message_receivers) = + out_message_mesh_builder.join(Role::Consumer).await.unwrap(); + let out_message_consumer_id = ConsumerId( + out_message_receivers + .consumer_id() + .unwrap() + .try_into() + .unwrap(), + ); + + ::log::info!("joined channels"); + + let connection_handles = Rc::new(RefCell::new(ConnectionHandles::default())); + + // Periodically clean connections + TimerActionRepeat::repeat_into( + enclose!((config, connection_handles, opt_tls_config) move || { + clean_connections( + config.clone(), + connection_handles.clone(), + server_start_instant, + opt_tls_config.clone(), + ) + }), + tq_prioritized, + ) + .unwrap(); + + for (_, out_message_receiver) in out_message_receivers.streams() { + spawn_local_into( + receive_out_messages(out_message_receiver, connection_handles.clone()), + tq_regular, + ) + .unwrap() + .detach(); + } + + let mut incoming = listener.incoming(); + + while let Some(stream) = incoming.next().await { + match stream { + Err(err) => { + ::log::error!("accept connection: {:#}", err); + } + Ok(stream) => { + let ip_version = match stream.peer_addr() { + Ok(addr) => IpVersion::canonical_from_ip(addr.ip()), + Err(err) => { + ::log::info!("could not extract ip version (v4 or v6): {:#}", err); + + continue; + } + }; + + let (out_message_sender, out_message_receiver) = new_bounded(LOCAL_CHANNEL_SIZE); + let out_message_sender = Rc::new(out_message_sender); + + let (close_conn_sender, close_conn_receiver) = new_bounded(1); + + let connection_valid_until = Rc::new(RefCell::new(ValidUntil::new( + server_start_instant, + config.cleaning.max_connection_idle, + ))); + + let connection_handle = ConnectionHandle { + close_conn_sender, + out_message_sender: out_message_sender.clone(), + valid_until: connection_valid_until.clone(), + opt_tls_config: opt_tls_config.as_ref().map(|c| c.load_full()), + valid_until_after_tls_update: None, + }; + + let connection_id = connection_handles.borrow_mut().insert(connection_handle); + + spawn_local_into( + enclose!(( + config, + access_list, + in_message_senders, + connection_valid_until, + opt_tls_config, + control_message_senders, + connection_handles + ) async move { + let runner = ConnectionRunner { + config, + access_list, + in_message_senders, + tq_prioritized, + tq_regular, + connection_valid_until, + out_message_sender, + out_message_receiver, + close_conn_receiver, + server_start_instant, + out_message_consumer_id, + connection_id, + opt_tls_config, + ip_version + }; + + runner.run(control_message_senders, stream).await; + + connection_handles.borrow_mut().remove(connection_id); + }), + tq_regular, + ) + .unwrap() + .detach(); + } + } + } +} + +async fn clean_connections( + config: Rc, + connection_slab: Rc>, + server_start_instant: ServerStartInstant, + opt_tls_config: Option>>, +) -> Option { + let now = server_start_instant.seconds_elapsed(); + let opt_current_tls_config = opt_tls_config.map(|c| c.load_full()); + + connection_slab.borrow_mut().retain(|_, reference| { + let mut keep = true; + + // Handle case when connection runs on an old TLS certificate + if let Some(valid_until) = reference.valid_until_after_tls_update { + if !valid_until.valid(now) { + keep = false; + } + } else if let Some(false) = opt_current_tls_config + .as_ref() + .zip(reference.opt_tls_config.as_ref()) + .map(|(a, b)| Arc::ptr_eq(a, b)) + { + reference.valid_until_after_tls_update = Some(ValidUntil::new( + server_start_instant, + config.cleaning.close_after_tls_update_grace_period, + )); + } + + keep &= reference.valid_until.borrow().valid(now); + + if keep { + true + } else { + if let Err(err) = reference.close_conn_sender.try_send(()) { + ::log::info!("couldn't tell connection to close: {:#}", err); + } + + false + } + }); + + ::log::info!( + "cleaned connections in worker {}, {} references remaining", + WORKER_INDEX.get(), + connection_slab.borrow_mut().len() + ); + + #[cfg(feature = "metrics")] + { + // Increment gauges by zero to prevent them from being removed due to + // idleness + + let worker_index = WORKER_INDEX.with(|index| index.get()).to_string(); + + if config.network.address.is_ipv4() || !config.network.only_ipv6 { + ::metrics::increment_gauge!( + "aquatic_active_connections", + 0.0, + "ip_version" => "4", + "worker_index" => worker_index.clone(), + ); + } + if config.network.address.is_ipv6() { + ::metrics::increment_gauge!( + "aquatic_active_connections", + 0.0, + "ip_version" => "6", + "worker_index" => worker_index, + ); + } + } + + Some(Duration::from_secs( + config.cleaning.connection_cleaning_interval, + )) +} + +async fn receive_out_messages( + mut out_message_receiver: ConnectedReceiver<(OutMessageMeta, OutMessage)>, + connection_references: Rc>, +) { + let connection_references = &connection_references; + + while let Some((meta, out_message)) = out_message_receiver.next().await { + if let Some(reference) = connection_references.borrow().get(meta.connection_id) { + match reference.out_message_sender.try_send((meta, out_message)) { + Ok(()) => {} + Err(GlommioError::Closed(_)) => {} + Err(GlommioError::WouldBlock(_)) => {} + Err(err) => { + ::log::debug!( + "Couldn't send out_message from shared channel to local receiver: {:?}", + err + ); + } + } + } + } +} + +fn create_tcp_listener( + config: &Config, + priv_dropper: PrivilegeDropper, +) -> anyhow::Result { + let domain = if config.network.address.is_ipv4() { + socket2::Domain::IPV4 + } else { + socket2::Domain::IPV6 + }; + + ::log::info!("creating socket.."); + + let socket = socket2::Socket::new(domain, socket2::Type::STREAM, Some(socket2::Protocol::TCP)) + .with_context(|| "create socket")?; + + if config.network.only_ipv6 { + ::log::info!("setting socket to ipv6 only.."); + + socket + .set_only_v6(true) + .with_context(|| "socket: set only ipv6")?; + } + + ::log::info!("setting SO_REUSEPORT on socket.."); + + socket + .set_reuse_port(true) + .with_context(|| "socket: set reuse port")?; + + ::log::info!("binding socket.."); + + socket + .bind(&config.network.address.into()) + .with_context(|| format!("socket: bind to {}", config.network.address))?; + + ::log::info!("listening on socket.."); + + socket + .listen(config.network.tcp_backlog) + .with_context(|| format!("socket: listen {}", config.network.address))?; + + ::log::info!("running PrivilegeDropper::after_socket_creation.."); + + priv_dropper.after_socket_creation()?; + + ::log::info!("casting socket to glommio TcpListener.."); + + Ok(unsafe { TcpListener::from_raw_fd(socket.into_raw_fd()) }) +} + +#[cfg(feature = "metrics")] +fn ip_version_to_metrics_str(ip_version: IpVersion) -> &'static str { + match ip_version { + IpVersion::V4 => "4", + IpVersion::V6 => "6", + } +} + +fn calculate_in_message_consumer_index(config: &Config, info_hash: InfoHash) -> usize { + (info_hash.0[0] as usize) % config.swarm_workers +} diff --git a/crates/ws/src/workers/swarm.rs b/crates/ws/src/workers/swarm.rs deleted file mode 100644 index 6cbf8dd..0000000 --- a/crates/ws/src/workers/swarm.rs +++ /dev/null @@ -1,541 +0,0 @@ -use std::cell::RefCell; -use std::rc::Rc; -use std::sync::Arc; -use std::time::Duration; - -use aquatic_common::access_list::{create_access_list_cache, AccessListArcSwap, AccessListCache}; -use futures::StreamExt; -use glommio::channels::channel_mesh::{MeshBuilder, Partial, Role, Senders}; -use glommio::enclose; -use glommio::prelude::*; -use glommio::timer::TimerActionRepeat; -use hashbrown::HashMap; -use rand::{rngs::SmallRng, SeedableRng}; - -use aquatic_common::{ - extract_response_peers, IndexMap, PanicSentinel, SecondsSinceServerStart, ServerStartInstant, -}; -use aquatic_ws_protocol::*; - -use crate::common::*; -use crate::config::Config; -use crate::SHARED_IN_CHANNEL_SIZE; - -#[cfg(feature = "metrics")] -thread_local! { static WORKER_INDEX: ::std::cell::Cell = Default::default() } - -#[derive(PartialEq, Eq, Clone, Copy, Debug)] -enum PeerStatus { - Seeding, - Leeching, - Stopped, -} - -impl PeerStatus { - /// Determine peer status from announce event and number of bytes left. - /// - /// Likely, the last branch will be taken most of the time. - #[inline] - fn from_event_and_bytes_left(event: AnnounceEvent, opt_bytes_left: Option) -> Self { - if let AnnounceEvent::Stopped = event { - Self::Stopped - } else if let Some(0) = opt_bytes_left { - Self::Seeding - } else { - Self::Leeching - } - } -} - -#[derive(Clone, Copy)] -struct Peer { - pub consumer_id: ConsumerId, - pub connection_id: ConnectionId, - pub seeder: bool, - pub valid_until: ValidUntil, -} - -type PeerMap = IndexMap; - -struct TorrentData { - pub peers: PeerMap, - pub num_seeders: usize, -} - -impl Default for TorrentData { - #[inline] - fn default() -> Self { - Self { - peers: Default::default(), - num_seeders: 0, - } - } -} - -impl TorrentData { - pub fn remove_peer(&mut self, peer_id: PeerId) { - if let Some(peer) = self.peers.remove(&peer_id) { - if peer.seeder { - self.num_seeders -= 1; - } - } - } - - pub fn num_leechers(&self) -> usize { - self.peers.len() - self.num_seeders - } -} - -type TorrentMap = IndexMap; - -#[derive(Default)] -struct TorrentMaps { - pub ipv4: TorrentMap, - pub ipv6: TorrentMap, -} - -impl TorrentMaps { - fn clean( - &mut self, - config: &Config, - access_list: &Arc, - server_start_instant: ServerStartInstant, - ) { - let mut access_list_cache = create_access_list_cache(access_list); - let now = server_start_instant.seconds_elapsed(); - - Self::clean_torrent_map(config, &mut access_list_cache, &mut self.ipv4, now, "4"); - Self::clean_torrent_map(config, &mut access_list_cache, &mut self.ipv6, now, "6"); - } - - fn clean_torrent_map( - config: &Config, - access_list_cache: &mut AccessListCache, - torrent_map: &mut TorrentMap, - now: SecondsSinceServerStart, - ip_version: &'static str, - ) { - let mut total_num_peers = 0u64; - - torrent_map.retain(|info_hash, torrent_data| { - if !access_list_cache - .load() - .allows(config.access_list.mode, &info_hash.0) - { - return false; - } - - let num_seeders = &mut torrent_data.num_seeders; - - torrent_data.peers.retain(|_, peer| { - let keep = peer.valid_until.valid(now); - - if (!keep) & peer.seeder { - *num_seeders -= 1; - } - - keep - }); - - total_num_peers += torrent_data.peers.len() as u64; - - !torrent_data.peers.is_empty() - }); - - torrent_map.shrink_to_fit(); - - let total_num_peers = total_num_peers as f64; - - #[cfg(feature = "metrics")] - ::metrics::gauge!( - "aquatic_peers", - total_num_peers, - "ip_version" => ip_version, - "worker_index" => WORKER_INDEX.with(|index| index.get()).to_string(), - ); - } -} - -pub async fn run_swarm_worker( - _sentinel: PanicSentinel, - config: Config, - state: State, - control_message_mesh_builder: MeshBuilder, - in_message_mesh_builder: MeshBuilder<(InMessageMeta, InMessage), Partial>, - out_message_mesh_builder: MeshBuilder<(OutMessageMeta, OutMessage), Partial>, - server_start_instant: ServerStartInstant, - worker_index: usize, -) { - #[cfg(feature = "metrics")] - WORKER_INDEX.with(|index| index.set(worker_index)); - - let (_, mut control_message_receivers) = control_message_mesh_builder - .join(Role::Consumer) - .await - .unwrap(); - - let (_, mut in_message_receivers) = in_message_mesh_builder.join(Role::Consumer).await.unwrap(); - let (out_message_senders, _) = out_message_mesh_builder.join(Role::Producer).await.unwrap(); - - let out_message_senders = Rc::new(out_message_senders); - - let torrents = Rc::new(RefCell::new(TorrentMaps::default())); - let access_list = state.access_list; - - // Periodically clean torrents - TimerActionRepeat::repeat(enclose!((config, torrents, access_list) move || { - enclose!((config, torrents, access_list) move || async move { - torrents.borrow_mut().clean(&config, &access_list, server_start_instant); - - Some(Duration::from_secs(config.cleaning.torrent_cleaning_interval)) - })() - })); - - // Periodically update torrent count metrics - #[cfg(feature = "metrics")] - TimerActionRepeat::repeat(enclose!((config, torrents) move || { - enclose!((config, torrents) move || async move { - let torrents = torrents.borrow_mut(); - - ::metrics::gauge!( - "aquatic_torrents", - torrents.ipv4.len() as f64, - "ip_version" => "4", - "worker_index" => worker_index.to_string(), - ); - ::metrics::gauge!( - "aquatic_torrents", - torrents.ipv6.len() as f64, - "ip_version" => "6", - "worker_index" => worker_index.to_string(), - ); - - Some(Duration::from_secs(config.metrics.torrent_count_update_interval)) - })() - })); - - let mut handles = Vec::new(); - - for (_, receiver) in control_message_receivers.streams() { - let handle = - spawn_local(handle_control_message_stream(torrents.clone(), receiver)).detach(); - - handles.push(handle); - } - - for (_, receiver) in in_message_receivers.streams() { - let handle = spawn_local(handle_request_stream( - config.clone(), - torrents.clone(), - server_start_instant, - out_message_senders.clone(), - receiver, - )) - .detach(); - - handles.push(handle); - } - - for handle in handles { - handle.await; - } -} - -async fn handle_control_message_stream(torrents: Rc>, mut stream: S) -where - S: futures_lite::Stream + ::std::marker::Unpin, -{ - while let Some(message) = stream.next().await { - match message { - SwarmControlMessage::ConnectionClosed { - info_hash, - peer_id, - ip_version, - } => { - ::log::debug!("Removing peer from torrents because connection was closed"); - - if let IpVersion::V4 = ip_version { - if let Some(torrent_data) = torrents.borrow_mut().ipv4.get_mut(&info_hash) { - torrent_data.remove_peer(peer_id); - } - } else { - if let Some(torrent_data) = torrents.borrow_mut().ipv6.get_mut(&info_hash) { - torrent_data.remove_peer(peer_id); - } - } - } - } - } -} - -async fn handle_request_stream( - config: Config, - torrents: Rc>, - server_start_instant: ServerStartInstant, - out_message_senders: Rc>, - stream: S, -) where - S: futures_lite::Stream + ::std::marker::Unpin, -{ - let rng = Rc::new(RefCell::new(SmallRng::from_entropy())); - - let max_peer_age = config.cleaning.max_peer_age; - let peer_valid_until = Rc::new(RefCell::new(ValidUntil::new( - server_start_instant, - max_peer_age, - ))); - - TimerActionRepeat::repeat(enclose!((peer_valid_until) move || { - enclose!((peer_valid_until) move || async move { - *peer_valid_until.borrow_mut() = ValidUntil::new(server_start_instant, max_peer_age); - - Some(Duration::from_secs(1)) - })() - })); - - let config = &config; - let torrents = &torrents; - let peer_valid_until = &peer_valid_until; - let rng = &rng; - let out_message_senders = &out_message_senders; - - stream - .for_each_concurrent( - SHARED_IN_CHANNEL_SIZE, - move |(meta, in_message)| async move { - let mut out_messages = Vec::new(); - - match in_message { - InMessage::AnnounceRequest(request) => handle_announce_request( - &config, - &mut rng.borrow_mut(), - &mut torrents.borrow_mut(), - &mut out_messages, - peer_valid_until.borrow().to_owned(), - meta, - request, - ), - InMessage::ScrapeRequest(request) => handle_scrape_request( - &config, - &mut torrents.borrow_mut(), - &mut out_messages, - meta, - request, - ), - }; - - for (meta, out_message) in out_messages.drain(..) { - out_message_senders - .send_to(meta.out_message_consumer_id.0 as usize, (meta, out_message)) - .await - .expect("failed sending out_message to socket worker"); - - ::log::debug!("swarm worker sent OutMessage to socket worker"); - } - }, - ) - .await; -} - -fn handle_announce_request( - config: &Config, - rng: &mut SmallRng, - torrent_maps: &mut TorrentMaps, - out_messages: &mut Vec<(OutMessageMeta, OutMessage)>, - valid_until: ValidUntil, - request_sender_meta: InMessageMeta, - request: AnnounceRequest, -) { - let (torrent_data, ip_version): (&mut TorrentData, &'static str) = - if let IpVersion::V4 = request_sender_meta.ip_version { - (torrent_maps.ipv4.entry(request.info_hash).or_default(), "4") - } else { - (torrent_maps.ipv6.entry(request.info_hash).or_default(), "6") - }; - - // If there is already a peer with this peer_id, check that connection id - // is same as that of request sender. Otherwise, ignore request. Since - // peers have access to each others peer_id's, they could send requests - // using them, causing all sorts of issues. - if let Some(previous_peer) = torrent_data.peers.get(&request.peer_id) { - if request_sender_meta.connection_id != previous_peer.connection_id { - return; - } - } - - ::log::trace!("received request from {:?}", request_sender_meta); - - // Insert/update/remove peer who sent this request - { - let peer_status = PeerStatus::from_event_and_bytes_left( - request.event.unwrap_or_default(), - request.bytes_left, - ); - - let opt_removed_peer = match peer_status { - PeerStatus::Leeching => { - let peer = Peer { - connection_id: request_sender_meta.connection_id, - consumer_id: request_sender_meta.out_message_consumer_id, - seeder: false, - valid_until, - }; - - torrent_data.peers.insert(request.peer_id, peer) - } - PeerStatus::Seeding => { - torrent_data.num_seeders += 1; - - let peer = Peer { - connection_id: request_sender_meta.connection_id, - consumer_id: request_sender_meta.out_message_consumer_id, - seeder: true, - valid_until, - }; - - torrent_data.peers.insert(request.peer_id, peer) - } - PeerStatus::Stopped => torrent_data.peers.remove(&request.peer_id), - }; - - if let Some(&Peer { seeder: true, .. }) = opt_removed_peer.as_ref() { - torrent_data.num_seeders -= 1; - } - - #[cfg(feature = "metrics")] - match peer_status { - PeerStatus::Stopped if opt_removed_peer.is_some() => { - ::metrics::decrement_gauge!( - "aquatic_peers", - 1.0, - "ip_version" => ip_version, - "worker_index" => WORKER_INDEX.with(|index| index.get()).to_string(), - ); - } - PeerStatus::Leeching | PeerStatus::Seeding if opt_removed_peer.is_none() => { - ::metrics::increment_gauge!( - "aquatic_peers", - 1.0, - "ip_version" => ip_version, - "worker_index" => WORKER_INDEX.with(|index| index.get()).to_string(), - ); - } - _ => {} - } - } - - // If peer sent offers, send them on to random peers - if let Some(offers) = request.offers { - // FIXME: config: also maybe check this when parsing request - let max_num_peers_to_take = offers.len().min(config.protocol.max_offers); - - #[inline] - fn f(peer: &Peer) -> Peer { - *peer - } - - let offer_receivers: Vec = extract_response_peers( - rng, - &torrent_data.peers, - max_num_peers_to_take, - request.peer_id, - f, - ); - - for (offer, offer_receiver) in offers.into_iter().zip(offer_receivers) { - let offer_out_message = OfferOutMessage { - action: AnnounceAction::Announce, - info_hash: request.info_hash, - peer_id: request.peer_id, - offer: offer.offer, - offer_id: offer.offer_id, - }; - - let meta = OutMessageMeta { - out_message_consumer_id: offer_receiver.consumer_id, - connection_id: offer_receiver.connection_id, - pending_scrape_id: None, - }; - - out_messages.push((meta, OutMessage::OfferOutMessage(offer_out_message))); - ::log::trace!("sending middleman offer to {:?}", meta); - } - } - - // If peer sent answer, send it on to relevant peer - if let (Some(answer), Some(answer_receiver_id), Some(offer_id)) = ( - request.answer, - request.answer_to_peer_id, - request.answer_offer_id, - ) { - if let Some(answer_receiver) = torrent_data.peers.get(&answer_receiver_id) { - let answer_out_message = AnswerOutMessage { - action: AnnounceAction::Announce, - peer_id: request.peer_id, - info_hash: request.info_hash, - answer, - offer_id, - }; - - let meta = OutMessageMeta { - out_message_consumer_id: answer_receiver.consumer_id, - connection_id: answer_receiver.connection_id, - pending_scrape_id: None, - }; - - out_messages.push((meta, OutMessage::AnswerOutMessage(answer_out_message))); - ::log::trace!("sending middleman answer to {:?}", meta); - } - } - - let out_message = OutMessage::AnnounceResponse(AnnounceResponse { - action: AnnounceAction::Announce, - info_hash: request.info_hash, - complete: torrent_data.num_seeders, - incomplete: torrent_data.num_leechers(), - announce_interval: config.protocol.peer_announce_interval, - }); - - out_messages.push((request_sender_meta.into(), out_message)); -} - -fn handle_scrape_request( - config: &Config, - torrent_maps: &mut TorrentMaps, - out_messages: &mut Vec<(OutMessageMeta, OutMessage)>, - meta: InMessageMeta, - request: ScrapeRequest, -) { - let info_hashes = if let Some(info_hashes) = request.info_hashes { - info_hashes.as_vec() - } else { - return; - }; - - let num_to_take = info_hashes.len().min(config.protocol.max_scrape_torrents); - - let mut out_message = ScrapeResponse { - action: ScrapeAction::Scrape, - files: HashMap::with_capacity(num_to_take), - }; - - let torrent_map: &mut TorrentMap = if let IpVersion::V4 = meta.ip_version { - &mut torrent_maps.ipv4 - } else { - &mut torrent_maps.ipv6 - }; - - for info_hash in info_hashes.into_iter().take(num_to_take) { - if let Some(torrent_data) = torrent_map.get(&info_hash) { - let stats = ScrapeStatistics { - complete: torrent_data.num_seeders, - downloaded: 0, // No implementation planned - incomplete: torrent_data.num_leechers(), - }; - - out_message.files.insert(info_hash, stats); - } - } - - out_messages.push((meta.into(), OutMessage::ScrapeResponse(out_message))); -} diff --git a/crates/ws/src/workers/swarm/mod.rs b/crates/ws/src/workers/swarm/mod.rs new file mode 100644 index 0000000..8fff337 --- /dev/null +++ b/crates/ws/src/workers/swarm/mod.rs @@ -0,0 +1,183 @@ +mod storage; + +use std::cell::RefCell; +use std::rc::Rc; +use std::time::Duration; + +use futures::StreamExt; +use glommio::channels::channel_mesh::{MeshBuilder, Partial, Role, Senders}; +use glommio::enclose; +use glommio::prelude::*; +use glommio::timer::TimerActionRepeat; +use rand::{rngs::SmallRng, SeedableRng}; + +use aquatic_common::{PanicSentinel, ServerStartInstant}; +use aquatic_ws_protocol::*; + +use crate::common::*; +use crate::config::Config; +use crate::SHARED_IN_CHANNEL_SIZE; + +use self::storage::TorrentMaps; + +#[cfg(feature = "metrics")] +thread_local! { static WORKER_INDEX: ::std::cell::Cell = Default::default() } + +pub async fn run_swarm_worker( + _sentinel: PanicSentinel, + config: Config, + state: State, + control_message_mesh_builder: MeshBuilder, + in_message_mesh_builder: MeshBuilder<(InMessageMeta, InMessage), Partial>, + out_message_mesh_builder: MeshBuilder<(OutMessageMeta, OutMessage), Partial>, + server_start_instant: ServerStartInstant, + worker_index: usize, +) { + #[cfg(feature = "metrics")] + WORKER_INDEX.with(|index| index.set(worker_index)); + + let (_, mut control_message_receivers) = control_message_mesh_builder + .join(Role::Consumer) + .await + .unwrap(); + + let (_, mut in_message_receivers) = in_message_mesh_builder.join(Role::Consumer).await.unwrap(); + let (out_message_senders, _) = out_message_mesh_builder.join(Role::Producer).await.unwrap(); + + let out_message_senders = Rc::new(out_message_senders); + + let torrents = Rc::new(RefCell::new(TorrentMaps::default())); + let access_list = state.access_list; + + // Periodically clean torrents + TimerActionRepeat::repeat(enclose!((config, torrents, access_list) move || { + enclose!((config, torrents, access_list) move || async move { + torrents.borrow_mut().clean(&config, &access_list, server_start_instant); + + Some(Duration::from_secs(config.cleaning.torrent_cleaning_interval)) + })() + })); + + // Periodically update torrent count metrics + #[cfg(feature = "metrics")] + TimerActionRepeat::repeat(enclose!((config, torrents) move || { + enclose!((config, torrents) move || async move { + torrents.borrow_mut().update_torrent_count_metrics(); + + Some(Duration::from_secs(config.metrics.torrent_count_update_interval)) + })() + })); + + let mut handles = Vec::new(); + + for (_, receiver) in control_message_receivers.streams() { + let handle = + spawn_local(handle_control_message_stream(torrents.clone(), receiver)).detach(); + + handles.push(handle); + } + + for (_, receiver) in in_message_receivers.streams() { + let handle = spawn_local(handle_request_stream( + config.clone(), + torrents.clone(), + server_start_instant, + out_message_senders.clone(), + receiver, + )) + .detach(); + + handles.push(handle); + } + + for handle in handles { + handle.await; + } +} + +async fn handle_control_message_stream(torrents: Rc>, mut stream: S) +where + S: futures_lite::Stream + ::std::marker::Unpin, +{ + while let Some(message) = stream.next().await { + match message { + SwarmControlMessage::ConnectionClosed { + info_hash, + peer_id, + ip_version, + } => { + torrents + .borrow_mut() + .handle_connection_closed(info_hash, peer_id, ip_version); + } + } + } +} + +async fn handle_request_stream( + config: Config, + torrents: Rc>, + server_start_instant: ServerStartInstant, + out_message_senders: Rc>, + stream: S, +) where + S: futures_lite::Stream + ::std::marker::Unpin, +{ + let rng = Rc::new(RefCell::new(SmallRng::from_entropy())); + + let max_peer_age = config.cleaning.max_peer_age; + let peer_valid_until = Rc::new(RefCell::new(ValidUntil::new( + server_start_instant, + max_peer_age, + ))); + + // Periodically update peer_valid_until + TimerActionRepeat::repeat(enclose!((peer_valid_until) move || { + enclose!((peer_valid_until) move || async move { + *peer_valid_until.borrow_mut() = ValidUntil::new(server_start_instant, max_peer_age); + + Some(Duration::from_secs(1)) + })() + })); + + let config = &config; + let torrents = &torrents; + let peer_valid_until = &peer_valid_until; + let rng = &rng; + let out_message_senders = &out_message_senders; + + stream + .for_each_concurrent( + SHARED_IN_CHANNEL_SIZE, + move |(meta, in_message)| async move { + let mut out_messages = Vec::new(); + + match in_message { + InMessage::AnnounceRequest(request) => { + torrents.borrow_mut().handle_announce_request( + &config, + &mut rng.borrow_mut(), + &mut out_messages, + server_start_instant, + peer_valid_until.borrow().to_owned(), + meta, + request, + ) + } + InMessage::ScrapeRequest(request) => torrents + .borrow_mut() + .handle_scrape_request(&config, &mut out_messages, meta, request), + }; + + for (meta, out_message) in out_messages { + out_message_senders + .send_to(meta.out_message_consumer_id.0 as usize, (meta, out_message)) + .await + .expect("failed sending out_message to socket worker"); + + ::log::debug!("swarm worker sent OutMessage to socket worker"); + } + }, + ) + .await; +} diff --git a/crates/ws/src/workers/swarm/storage.rs b/crates/ws/src/workers/swarm/storage.rs new file mode 100644 index 0000000..1ec9d08 --- /dev/null +++ b/crates/ws/src/workers/swarm/storage.rs @@ -0,0 +1,488 @@ +use std::sync::Arc; + +use aquatic_common::access_list::{create_access_list_cache, AccessListArcSwap, AccessListCache}; +use hashbrown::HashMap; +use rand::rngs::SmallRng; + +use aquatic_common::{ + extract_response_peers, IndexMap, SecondsSinceServerStart, ServerStartInstant, +}; +use aquatic_ws_protocol::*; + +use crate::common::*; +use crate::config::Config; +use crate::workers::swarm::WORKER_INDEX; + +type TorrentMap = IndexMap; +type PeerMap = IndexMap; + +#[derive(Default)] +pub struct TorrentMaps { + ipv4: TorrentMap, + ipv6: TorrentMap, +} + +impl TorrentMaps { + pub fn handle_announce_request( + &mut self, + config: &Config, + rng: &mut SmallRng, + out_messages: &mut Vec<(OutMessageMeta, OutMessage)>, + server_start_instant: ServerStartInstant, + valid_until: ValidUntil, + request_sender_meta: InMessageMeta, + request: AnnounceRequest, + ) { + let (torrent_data, ip_version): (&mut TorrentData, &'static str) = + if let IpVersion::V4 = request_sender_meta.ip_version { + (self.ipv4.entry(request.info_hash).or_default(), "4") + } else { + (self.ipv6.entry(request.info_hash).or_default(), "6") + }; + + // If there is already a peer with this peer_id, check that connection id + // is same as that of request sender. Otherwise, ignore request. Since + // peers have access to each others peer_id's, they could send requests + // using them, causing all sorts of issues. + if let Some(previous_peer) = torrent_data.peers.get(&request.peer_id) { + if request_sender_meta.connection_id != previous_peer.connection_id { + return; + } + } + + ::log::trace!("received request from {:?}", request_sender_meta); + + // Insert/update/remove peer who sent this request + { + let peer_status = PeerStatus::from_event_and_bytes_left( + request.event.unwrap_or_default(), + request.bytes_left, + ); + + match torrent_data.peers.entry(request.peer_id) { + ::indexmap::map::Entry::Occupied(mut entry) => match peer_status { + PeerStatus::Leeching => { + let peer = entry.get_mut(); + + if peer.seeder { + torrent_data.num_seeders -= 1; + } + + peer.seeder = false; + peer.valid_until = valid_until; + } + PeerStatus::Seeding => { + let peer = entry.get_mut(); + + if !peer.seeder { + torrent_data.num_seeders += 1; + } + + peer.seeder = true; + peer.valid_until = valid_until; + } + PeerStatus::Stopped => { + let peer = entry.remove(); + + if peer.seeder { + torrent_data.num_seeders -= 1; + } + + #[cfg(feature = "metrics")] + ::metrics::decrement_gauge!( + "aquatic_peers", + 1.0, + "ip_version" => ip_version, + "worker_index" => WORKER_INDEX.with(|index| index.get()).to_string(), + ); + + return; + } + }, + ::indexmap::map::Entry::Vacant(entry) => match peer_status { + PeerStatus::Leeching => { + let peer = Peer { + connection_id: request_sender_meta.connection_id, + consumer_id: request_sender_meta.out_message_consumer_id, + seeder: false, + valid_until, + expecting_answers: Default::default(), + }; + + entry.insert(peer); + + #[cfg(feature = "metrics")] + ::metrics::increment_gauge!( + "aquatic_peers", + 1.0, + "ip_version" => ip_version, + "worker_index" => WORKER_INDEX.with(|index| index.get()).to_string(), + ); + } + PeerStatus::Seeding => { + torrent_data.num_seeders += 1; + + let peer = Peer { + connection_id: request_sender_meta.connection_id, + consumer_id: request_sender_meta.out_message_consumer_id, + seeder: true, + valid_until, + expecting_answers: Default::default(), + }; + + entry.insert(peer); + + #[cfg(feature = "metrics")] + ::metrics::increment_gauge!( + "aquatic_peers", + 1.0, + "ip_version" => ip_version, + "worker_index" => WORKER_INDEX.with(|index| index.get()).to_string(), + ); + } + PeerStatus::Stopped => return, + }, + } + }; + + // If peer sent offers, send them on to random peers + if let Some(offers) = request.offers { + // FIXME: config: also maybe check this when parsing request + let max_num_peers_to_take = offers.len().min(config.protocol.max_offers); + + #[inline] + fn convert_offer_receiver_peer( + peer_id: &PeerId, + peer: &Peer, + ) -> (PeerId, ConnectionId, ConsumerId) { + (*peer_id, peer.connection_id, peer.consumer_id) + } + + let offer_receivers: Vec<(PeerId, ConnectionId, ConsumerId)> = extract_response_peers( + rng, + &torrent_data.peers, + max_num_peers_to_take, + request.peer_id, + convert_offer_receiver_peer, + ); + + if let Some(peer) = torrent_data.peers.get_mut(&request.peer_id) { + for ( + offer, + ( + offer_receiver_peer_id, + offer_receiver_connection_id, + offer_receiver_consumer_id, + ), + ) in offers.into_iter().zip(offer_receivers) + { + let offer_out_message = OfferOutMessage { + action: AnnounceAction::Announce, + info_hash: request.info_hash, + peer_id: request.peer_id, + offer: offer.offer, + offer_id: offer.offer_id, + }; + + let meta = OutMessageMeta { + out_message_consumer_id: offer_receiver_consumer_id, + connection_id: offer_receiver_connection_id, + pending_scrape_id: None, + }; + + out_messages.push((meta, OutMessage::OfferOutMessage(offer_out_message))); + ::log::trace!("sending OfferOutMessage to {:?}", meta); + + peer.expecting_answers.insert( + ExpectingAnswer { + from_peer_id: offer_receiver_peer_id, + regarding_offer_id: offer.offer_id, + }, + ValidUntil::new(server_start_instant, config.cleaning.max_offer_age), + ); + } + } + } + + // If peer sent answer, send it on to relevant peer + if let (Some(answer), Some(answer_receiver_id), Some(offer_id)) = ( + request.answer, + request.answer_to_peer_id, + request.answer_offer_id, + ) { + if let Some(answer_receiver) = torrent_data.peers.get_mut(&answer_receiver_id) { + let expecting_answer = ExpectingAnswer { + from_peer_id: request.peer_id, + regarding_offer_id: offer_id, + }; + + if let Some(_) = answer_receiver.expecting_answers.remove(&expecting_answer) { + let answer_out_message = AnswerOutMessage { + action: AnnounceAction::Announce, + peer_id: request.peer_id, + info_hash: request.info_hash, + answer, + offer_id, + }; + + let meta = OutMessageMeta { + out_message_consumer_id: answer_receiver.consumer_id, + connection_id: answer_receiver.connection_id, + pending_scrape_id: None, + }; + + out_messages.push((meta, OutMessage::AnswerOutMessage(answer_out_message))); + ::log::trace!("sending AnswerOutMessage to {:?}", meta); + } else { + let error_message = ErrorResponse { + action: Some(ErrorResponseAction::Announce), + info_hash: Some(request.info_hash), + failure_reason: + "Could not find the offer corresponding to your answer. It may have expired." + .into(), + }; + + out_messages.push(( + request_sender_meta.into(), + OutMessage::ErrorResponse(error_message), + )); + } + } + } + + let out_message = OutMessage::AnnounceResponse(AnnounceResponse { + action: AnnounceAction::Announce, + info_hash: request.info_hash, + complete: torrent_data.num_seeders, + incomplete: torrent_data.num_leechers(), + announce_interval: config.protocol.peer_announce_interval, + }); + + out_messages.push((request_sender_meta.into(), out_message)); + } + + pub fn handle_scrape_request( + &mut self, + config: &Config, + out_messages: &mut Vec<(OutMessageMeta, OutMessage)>, + meta: InMessageMeta, + request: ScrapeRequest, + ) { + let info_hashes = if let Some(info_hashes) = request.info_hashes { + info_hashes.as_vec() + } else { + return; + }; + + let num_to_take = info_hashes.len().min(config.protocol.max_scrape_torrents); + + let mut out_message = ScrapeResponse { + action: ScrapeAction::Scrape, + files: HashMap::with_capacity(num_to_take), + }; + + let torrent_map: &mut TorrentMap = if let IpVersion::V4 = meta.ip_version { + &mut self.ipv4 + } else { + &mut self.ipv6 + }; + + for info_hash in info_hashes.into_iter().take(num_to_take) { + if let Some(torrent_data) = torrent_map.get(&info_hash) { + let stats = ScrapeStatistics { + complete: torrent_data.num_seeders, + downloaded: 0, // No implementation planned + incomplete: torrent_data.num_leechers(), + }; + + out_message.files.insert(info_hash, stats); + } + } + + out_messages.push((meta.into(), OutMessage::ScrapeResponse(out_message))); + } + + pub fn clean( + &mut self, + config: &Config, + access_list: &Arc, + server_start_instant: ServerStartInstant, + ) { + let mut access_list_cache = create_access_list_cache(access_list); + let now = server_start_instant.seconds_elapsed(); + + Self::clean_torrent_map(config, &mut access_list_cache, &mut self.ipv4, now, "4"); + Self::clean_torrent_map(config, &mut access_list_cache, &mut self.ipv6, now, "6"); + } + + fn clean_torrent_map( + config: &Config, + access_list_cache: &mut AccessListCache, + torrent_map: &mut TorrentMap, + now: SecondsSinceServerStart, + ip_version: &'static str, + ) { + let mut total_num_peers = 0u64; + + torrent_map.retain(|info_hash, torrent_data| { + if !access_list_cache + .load() + .allows(config.access_list.mode, &info_hash.0) + { + return false; + } + + let num_seeders = &mut torrent_data.num_seeders; + + torrent_data.peers.retain(|_, peer| { + peer.expecting_answers + .retain(|_, valid_until| valid_until.valid(now)); + peer.expecting_answers.shrink_to_fit(); + + let keep = peer.valid_until.valid(now); + + if (!keep) & peer.seeder { + *num_seeders -= 1; + } + + keep + }); + + total_num_peers += torrent_data.peers.len() as u64; + + torrent_data.peers.shrink_to_fit(); + + !torrent_data.peers.is_empty() + }); + + torrent_map.shrink_to_fit(); + + let total_num_peers = total_num_peers as f64; + + #[cfg(feature = "metrics")] + ::metrics::gauge!( + "aquatic_peers", + total_num_peers, + "ip_version" => ip_version, + "worker_index" => WORKER_INDEX.with(|index| index.get()).to_string(), + ); + } + + #[cfg(feature = "metrics")] + pub fn update_torrent_count_metrics(&self) { + ::metrics::gauge!( + "aquatic_torrents", + self.ipv4.len() as f64, + "ip_version" => "4", + "worker_index" => WORKER_INDEX.with(|index| index.get()).to_string(), + ); + ::metrics::gauge!( + "aquatic_torrents", + self.ipv6.len() as f64, + "ip_version" => "6", + "worker_index" => WORKER_INDEX.with(|index| index.get()).to_string(), + ); + } + + pub fn handle_connection_closed( + &mut self, + info_hash: InfoHash, + peer_id: PeerId, + ip_version: IpVersion, + ) { + ::log::debug!("Removing peer from torrents because connection was closed"); + + if let IpVersion::V4 = ip_version { + if let Some(torrent_data) = self.ipv4.get_mut(&info_hash) { + torrent_data.remove_peer(peer_id); + + #[cfg(feature = "metrics")] + ::metrics::decrement_gauge!( + "aquatic_peers", + 1.0, + "ip_version" => "4", + "worker_index" => WORKER_INDEX.with(|index| index.get()).to_string(), + ); + } + } else { + if let Some(torrent_data) = self.ipv6.get_mut(&info_hash) { + torrent_data.remove_peer(peer_id); + + #[cfg(feature = "metrics")] + ::metrics::decrement_gauge!( + "aquatic_peers", + 1.0, + "ip_version" => "6", + "worker_index" => WORKER_INDEX.with(|index| index.get()).to_string(), + ); + } + } + } +} + +struct TorrentData { + peers: PeerMap, + num_seeders: usize, +} + +impl Default for TorrentData { + #[inline] + fn default() -> Self { + Self { + peers: Default::default(), + num_seeders: 0, + } + } +} + +impl TorrentData { + fn remove_peer(&mut self, peer_id: PeerId) { + if let Some(peer) = self.peers.remove(&peer_id) { + if peer.seeder { + self.num_seeders -= 1; + } + } + } + + fn num_leechers(&self) -> usize { + self.peers.len() - self.num_seeders + } +} + +#[derive(Clone, Debug)] +struct Peer { + pub consumer_id: ConsumerId, + pub connection_id: ConnectionId, + pub seeder: bool, + pub valid_until: ValidUntil, + pub expecting_answers: IndexMap, +} + +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct ExpectingAnswer { + pub from_peer_id: PeerId, + pub regarding_offer_id: OfferId, +} + +#[derive(PartialEq, Eq, Clone, Copy, Debug)] +enum PeerStatus { + Seeding, + Leeching, + Stopped, +} + +impl PeerStatus { + /// Determine peer status from announce event and number of bytes left. + /// + /// Likely, the last branch will be taken most of the time. + #[inline] + fn from_event_and_bytes_left(event: AnnounceEvent, opt_bytes_left: Option) -> Self { + if let AnnounceEvent::Stopped = event { + Self::Stopped + } else if let Some(0) = opt_bytes_left { + Self::Seeding + } else { + Self::Leeching + } + } +} diff --git a/crates/ws_load_test/src/network.rs b/crates/ws_load_test/src/network.rs index cafa52b..2a03ad2 100644 --- a/crates/ws_load_test/src/network.rs +++ b/crates/ws_load_test/src/network.rs @@ -125,6 +125,7 @@ impl Connection { &self.load_test_state, &mut self.rng, self.peer_id, + self.send_answer.is_none(), ); // If self.send_answer is set and request is announce request, make diff --git a/crates/ws_load_test/src/utils.rs b/crates/ws_load_test/src/utils.rs index 354e180..bb17443 100644 --- a/crates/ws_load_test/src/utils.rs +++ b/crates/ws_load_test/src/utils.rs @@ -12,6 +12,7 @@ pub fn create_random_request( state: &LoadTestState, rng: &mut impl Rng, peer_id: PeerId, + announce_gen_offers: bool, ) -> InMessage { let weights = [ config.torrents.weight_announce as u32, @@ -23,7 +24,9 @@ pub fn create_random_request( let dist = WeightedIndex::new(&weights).expect("random request weighted index"); match items[dist.sample(rng)] { - RequestType::Announce => create_announce_request(config, state, rng, peer_id), + RequestType::Announce => { + create_announce_request(config, state, rng, peer_id, announce_gen_offers) + } RequestType::Scrape => create_scrape_request(config, state, rng), } } @@ -34,6 +37,7 @@ fn create_announce_request( state: &LoadTestState, rng: &mut impl Rng, peer_id: PeerId, + gen_offers: bool, ) -> InMessage { let (event, bytes_left) = { if rng.gen_bool(config.torrents.peer_seeder_probability) { @@ -45,17 +49,23 @@ fn create_announce_request( let info_hash_index = select_info_hash_index(config, &state, rng); - let mut offers = Vec::with_capacity(config.torrents.offers_per_request); + let offers = if gen_offers { + let mut offers = Vec::with_capacity(config.torrents.offers_per_request); - for _ in 0..config.torrents.offers_per_request { - offers.push(AnnounceRequestOffer { - offer_id: OfferId(rng.gen()), - offer: RtcOffer { - t: RtcOfferType::Offer, - sdp: "abcdefg-abcdefg-abcdefg-abcdefg-abcdefg-abcdefg-abcdefg-abcdefg-abcdefg-abcdefg-abcdefg-abcdefg-abcdefg-abcdefg-abcdefg-".into() - }, - }) - } + for _ in 0..config.torrents.offers_per_request { + offers.push(AnnounceRequestOffer { + offer_id: OfferId(rng.gen()), + offer: RtcOffer { + t: RtcOfferType::Offer, + sdp: "abcdefg-abcdefg-abcdefg-abcdefg-abcdefg-abcdefg-abcdefg-abcdefg-abcdefg-abcdefg-abcdefg-abcdefg-abcdefg-abcdefg-abcdefg-".into() + }, + }) + } + + offers + } else { + Vec::new() + }; InMessage::AnnounceRequest(AnnounceRequest { action: AnnounceAction::Announce, diff --git a/crates/ws_protocol/src/common.rs b/crates/ws_protocol/src/common.rs index 08e00f8..f465f04 100644 --- a/crates/ws_protocol/src/common.rs +++ b/crates/ws_protocol/src/common.rs @@ -20,7 +20,7 @@ pub struct InfoHash( pub [u8; 20], ); -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq, Serialize, Deserialize)] #[serde(transparent)] pub struct OfferId( #[serde(