Move all crates to new crates dir

crates/ws/src/common.rs

@@ -0,0 +1,76 @@
+use std::{net::IpAddr, sync::Arc};
+
+use aquatic_common::access_list::AccessListArcSwap;
+
+pub use aquatic_common::ValidUntil;
+use aquatic_ws_protocol::{InfoHash, PeerId};
+
+#[derive(Copy, Clone, Debug)]
+pub enum IpVersion {
+    V4,
+    V6,
+}
+
+impl IpVersion {
+    pub fn canonical_from_ip(ip: IpAddr) -> IpVersion {
+        match ip {
+            IpAddr::V4(_) => Self::V4,
+            IpAddr::V6(addr) => match addr.octets() {
+                [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0xff, 0xff, _, _, _, _] => Self::V4,
+                _ => Self::V6,
+            },
+        }
+    }
+}
+
+#[derive(Default, Clone)]
+pub struct State {
+    pub access_list: Arc<AccessListArcSwap>,
+}
+
+#[derive(Copy, Clone, Debug)]
+pub struct PendingScrapeId(pub u8);
+
+#[derive(Copy, Clone, Debug)]
+pub struct ConsumerId(pub u8);
+
+#[derive(Clone, Copy, Debug, PartialEq)]
+pub struct ConnectionId(pub usize);
+
+#[derive(Clone, Copy, Debug)]
+pub struct InMessageMeta {
+    /// Index of socket worker responsible for this connection. Required for
+    /// sending back response through correct channel to correct worker.
+    pub out_message_consumer_id: ConsumerId,
+    pub connection_id: ConnectionId,
+    pub ip_version: IpVersion,
+    pub pending_scrape_id: Option<PendingScrapeId>,
+}
+
+#[derive(Clone, Copy, Debug)]
+pub struct OutMessageMeta {
+    /// Index of socket worker responsible for this connection. Required for
+    /// sending back response through correct channel to correct worker.
+    pub out_message_consumer_id: ConsumerId,
+    pub connection_id: ConnectionId,
+    pub pending_scrape_id: Option<PendingScrapeId>,
+}
+
+impl Into<OutMessageMeta> for InMessageMeta {
+    fn into(self) -> OutMessageMeta {
+        OutMessageMeta {
+            out_message_consumer_id: self.out_message_consumer_id,
+            connection_id: self.connection_id,
+            pending_scrape_id: self.pending_scrape_id,
+        }
+    }
+}
+
+#[derive(Clone, Copy, Debug)]
+pub enum SwarmControlMessage {
+    ConnectionClosed {
+        info_hash: InfoHash,
+        peer_id: PeerId,
+        ip_version: IpVersion,
+    },
+}

crates/ws/src/config.rs

@@ -0,0 +1,191 @@
+use std::net::SocketAddr;
+use std::path::PathBuf;
+
+use aquatic_common::cpu_pinning::asc::CpuPinningConfigAsc;
+use aquatic_common::{access_list::AccessListConfig, privileges::PrivilegeConfig};
+use serde::Deserialize;
+
+use aquatic_common::cli::LogLevel;
+use aquatic_toml_config::TomlConfig;
+
+/// aquatic_ws configuration
+///
+/// Running behind a reverse proxy is supported, but IPv4 peer requests have
+/// to be proxied to IPv4 requests, and IPv6 requests to IPv6 requests.
+#[derive(Clone, Debug, PartialEq, TomlConfig, Deserialize)]
+#[serde(default, deny_unknown_fields)]
+pub struct Config {
+    /// Socket workers receive requests from the socket, parse them and send
+    /// them on to the swarm workers. They then receive responses from the
+    /// swarm workers, encode them and send them back over the socket.
+    pub socket_workers: usize,
+    /// Swarm workers receive a number of requests from socket workers,
+    /// generate responses and send them back to the socket workers.
+    pub swarm_workers: usize,
+    pub log_level: LogLevel,
+    pub network: NetworkConfig,
+    pub protocol: ProtocolConfig,
+    pub cleaning: CleaningConfig,
+    pub privileges: PrivilegeConfig,
+    pub access_list: AccessListConfig,
+    #[cfg(feature = "metrics")]
+    pub metrics: MetricsConfig,
+    pub cpu_pinning: CpuPinningConfigAsc,
+}
+
+impl Default for Config {
+    fn default() -> Self {
+        Self {
+            socket_workers: 1,
+            swarm_workers: 1,
+            log_level: LogLevel::default(),
+            network: NetworkConfig::default(),
+            protocol: ProtocolConfig::default(),
+            cleaning: CleaningConfig::default(),
+            privileges: PrivilegeConfig::default(),
+            access_list: AccessListConfig::default(),
+            #[cfg(feature = "metrics")]
+            metrics: Default::default(),
+            cpu_pinning: Default::default(),
+        }
+    }
+}
+
+impl aquatic_common::cli::Config for Config {
+    fn get_log_level(&self) -> Option<LogLevel> {
+        Some(self.log_level)
+    }
+}
+
+#[derive(Clone, Debug, PartialEq, TomlConfig, Deserialize)]
+#[serde(default, deny_unknown_fields)]
+pub struct NetworkConfig {
+    /// Bind to this address
+    pub address: SocketAddr,
+    /// Only allow access over IPv6
+    pub only_ipv6: bool,
+    /// Maximum number of pending TCP connections
+    pub tcp_backlog: i32,
+
+    /// Enable TLS
+    pub enable_tls: bool,
+    /// Path to TLS certificate (DER-encoded X.509)
+    pub tls_certificate_path: PathBuf,
+    /// Path to TLS private key (DER-encoded ASN.1 in PKCS#8 or PKCS#1 format)
+    pub tls_private_key_path: PathBuf,
+
+    pub websocket_max_message_size: usize,
+    pub websocket_max_frame_size: usize,
+    pub websocket_write_buffer_size: usize,
+
+    /// Return a HTTP 200 Ok response when receiving GET /health. Can not be
+    /// combined with enable_tls.
+    pub enable_http_health_checks: bool,
+}
+
+impl Default for NetworkConfig {
+    fn default() -> Self {
+        Self {
+            address: SocketAddr::from(([0, 0, 0, 0], 3000)),
+            only_ipv6: false,
+            tcp_backlog: 1024,
+
+            enable_tls: false,
+            tls_certificate_path: "".into(),
+            tls_private_key_path: "".into(),
+
+            websocket_max_message_size: 64 * 1024,
+            websocket_max_frame_size: 16 * 1024,
+            websocket_write_buffer_size: 8 * 1024,
+
+            enable_http_health_checks: false,
+        }
+    }
+}
+
+#[derive(Clone, Debug, PartialEq, TomlConfig, Deserialize)]
+#[serde(default, deny_unknown_fields)]
+pub struct ProtocolConfig {
+    /// Maximum number of torrents to accept in scrape request
+    pub max_scrape_torrents: usize,
+    /// Maximum number of offers to accept in announce request
+    pub max_offers: usize,
+    /// Ask peers to announce this often (seconds)
+    pub peer_announce_interval: usize,
+}
+
+impl Default for ProtocolConfig {
+    fn default() -> Self {
+        Self {
+            max_scrape_torrents: 255,
+            max_offers: 10,
+            peer_announce_interval: 120,
+        }
+    }
+}
+
+#[derive(Clone, Debug, PartialEq, TomlConfig, Deserialize)]
+#[serde(default, deny_unknown_fields)]
+pub struct CleaningConfig {
+    /// Clean peers this often (seconds)
+    pub torrent_cleaning_interval: u64,
+    /// Remove peers that have not announced for this long (seconds)
+    pub max_peer_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)
+    pub max_connection_idle: u32,
+}
+
+impl Default for CleaningConfig {
+    fn default() -> Self {
+        Self {
+            torrent_cleaning_interval: 30,
+            max_peer_age: 1800,
+            max_connection_idle: 60 * 5,
+            connection_cleaning_interval: 30,
+        }
+    }
+}
+
+#[cfg(feature = "metrics")]
+#[derive(Clone, Debug, PartialEq, TomlConfig, Deserialize)]
+#[serde(default, deny_unknown_fields)]
+pub struct MetricsConfig {
+    /// Run a prometheus endpoint
+    pub run_prometheus_endpoint: bool,
+    /// Address to run prometheus endpoint on
+    pub prometheus_endpoint_address: SocketAddr,
+    /// Update metrics for torrent count this often (seconds)
+    pub torrent_count_update_interval: u64,
+    /// Serve information on peer clients
+    ///
+    /// Expect a certain CPU hit
+    pub peer_clients: bool,
+    /// Serve information on all peer id prefixes
+    ///
+    /// Requires `peer_clients` to be activated.
+    ///
+    /// Expect a certain CPU hit
+    pub peer_id_prefixes: bool,
+}
+
+#[cfg(feature = "metrics")]
+impl Default for MetricsConfig {
+    fn default() -> Self {
+        Self {
+            run_prometheus_endpoint: false,
+            prometheus_endpoint_address: SocketAddr::from(([0, 0, 0, 0], 9000)),
+            torrent_count_update_interval: 10,
+            peer_clients: false,
+            peer_id_prefixes: false,
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::Config;
+
+    ::aquatic_toml_config::gen_serialize_deserialize_test!(Config);
+}

crates/ws/src/lib.rs

@@ -0,0 +1,197 @@
+pub mod common;
+pub mod config;
+pub mod workers;
+
+use std::sync::Arc;
+use std::time::Duration;
+
+use anyhow::Context;
+use aquatic_common::cpu_pinning::glommio::{get_worker_placement, set_affinity_for_util_worker};
+use aquatic_common::cpu_pinning::WorkerIndex;
+use aquatic_common::rustls_config::create_rustls_config;
+use aquatic_common::{PanicSentinelWatcher, ServerStartInstant};
+use glommio::{channels::channel_mesh::MeshBuilder, prelude::*};
+use signal_hook::{
+    consts::{SIGTERM, SIGUSR1},
+    iterator::Signals,
+};
+
+use aquatic_common::access_list::update_access_list;
+use aquatic_common::privileges::PrivilegeDropper;
+
+use common::*;
+use config::Config;
+
+pub const APP_NAME: &str = "aquatic_ws: WebTorrent tracker";
+pub const APP_VERSION: &str = env!("CARGO_PKG_VERSION");
+
+pub const SHARED_IN_CHANNEL_SIZE: usize = 1024;
+
+pub fn run(config: Config) -> ::anyhow::Result<()> {
+    if config.network.enable_tls && config.network.enable_http_health_checks {
+        return Err(anyhow::anyhow!(
+            "configuration: network.enable_tls and network.enable_http_health_check can't both be set to true"
+        ));
+    }
+
+    let mut signals = Signals::new([SIGUSR1, SIGTERM])?;
+
+    #[cfg(feature = "prometheus")]
+    if config.metrics.run_prometheus_endpoint {
+        use metrics_exporter_prometheus::PrometheusBuilder;
+
+        let idle_timeout = config
+            .cleaning
+            .connection_cleaning_interval
+            .max(config.cleaning.torrent_cleaning_interval)
+            .max(config.metrics.torrent_count_update_interval)
+            * 2;
+
+        PrometheusBuilder::new()
+            .idle_timeout(
+                metrics_util::MetricKindMask::GAUGE,
+                Some(Duration::from_secs(idle_timeout)),
+            )
+            .with_http_listener(config.metrics.prometheus_endpoint_address)
+            .install()
+            .with_context(|| {
+                format!(
+                    "Install prometheus endpoint on {}",
+                    config.metrics.prometheus_endpoint_address
+                )
+            })?;
+    }
+
+    let state = State::default();
+
+    update_access_list(&config.access_list, &state.access_list)?;
+
+    let num_peers = config.socket_workers + config.swarm_workers;
+
+    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 (sentinel_watcher, sentinel) = PanicSentinelWatcher::create_with_sentinel();
+    let priv_dropper = PrivilegeDropper::new(config.privileges.clone(), config.socket_workers);
+
+    let opt_tls_config = if config.network.enable_tls {
+        Some(Arc::new(
+            create_rustls_config(
+                &config.network.tls_certificate_path,
+                &config.network.tls_private_key_path,
+            )
+            .with_context(|| "create rustls config")?,
+        ))
+    } else {
+        None
+    };
+
+    let server_start_instant = ServerStartInstant::new();
+
+    let mut executors = Vec::new();
+
+    for i in 0..(config.socket_workers) {
+        let sentinel = sentinel.clone();
+        let config = config.clone();
+        let state = state.clone();
+        let opt_tls_config = opt_tls_config.clone();
+        let control_mesh_builder = control_mesh_builder.clone();
+        let request_mesh_builder = request_mesh_builder.clone();
+        let response_mesh_builder = response_mesh_builder.clone();
+        let priv_dropper = priv_dropper.clone();
+
+        let placement = get_worker_placement(
+            &config.cpu_pinning,
+            config.socket_workers,
+            config.swarm_workers,
+            WorkerIndex::SocketWorker(i),
+        )?;
+        let builder = LocalExecutorBuilder::new(placement).name(&format!("socket-{:02}", i + 1));
+
+        let executor = builder
+            .spawn(move || async move {
+                workers::socket::run_socket_worker(
+                    sentinel,
+                    config,
+                    state,
+                    opt_tls_config,
+                    control_mesh_builder,
+                    request_mesh_builder,
+                    response_mesh_builder,
+                    priv_dropper,
+                    server_start_instant,
+                    i,
+                )
+                .await
+            })
+            .map_err(|err| anyhow::anyhow!("Spawning executor failed: {:#}", err))?;
+
+        executors.push(executor);
+    }
+
+    ::log::info!("spawned socket workers");
+
+    for i in 0..(config.swarm_workers) {
+        let sentinel = sentinel.clone();
+        let config = config.clone();
+        let state = state.clone();
+        let control_mesh_builder = control_mesh_builder.clone();
+        let request_mesh_builder = request_mesh_builder.clone();
+        let response_mesh_builder = response_mesh_builder.clone();
+
+        let placement = get_worker_placement(
+            &config.cpu_pinning,
+            config.socket_workers,
+            config.swarm_workers,
+            WorkerIndex::SwarmWorker(i),
+        )?;
+        let builder = LocalExecutorBuilder::new(placement).name(&format!("swarm-{:02}", i + 1));
+
+        let executor = builder
+            .spawn(move || async move {
+                workers::swarm::run_swarm_worker(
+                    sentinel,
+                    config,
+                    state,
+                    control_mesh_builder,
+                    request_mesh_builder,
+                    response_mesh_builder,
+                    server_start_instant,
+                    i,
+                )
+                .await
+            })
+            .map_err(|err| anyhow::anyhow!("Spawning executor failed: {:#}", err))?;
+
+        executors.push(executor);
+    }
+
+    ::log::info!("spawned swarm workers");
+
+    if config.cpu_pinning.active {
+        set_affinity_for_util_worker(
+            &config.cpu_pinning,
+            config.socket_workers,
+            config.swarm_workers,
+        )?;
+    }
+
+    for signal in &mut signals {
+        match signal {
+            SIGUSR1 => {
+                let _ = update_access_list(&config.access_list, &state.access_list);
+            }
+            SIGTERM => {
+                if sentinel_watcher.panic_was_triggered() {
+                    return Err(anyhow::anyhow!("worker thread panicked"));
+                } else {
+                    return Ok(());
+                }
+            }
+            _ => unreachable!(),
+        }
+    }
+
+    Ok(())
+}

crates/ws/src/main.rs

@@ -0,0 +1,14 @@
+use aquatic_common::cli::run_app_with_cli_and_config;
+use aquatic_ws::config::Config;
+
+#[global_allocator]
+static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
+
+fn main() {
+    run_app_with_cli_and_config::<Config>(
+        aquatic_ws::APP_NAME,
+        aquatic_ws::APP_VERSION,
+        aquatic_ws::run,
+        None,
+    )
+}

crates/ws/src/workers/mod.rs

@@ -0,0 +1,2 @@
+pub mod socket;
+pub mod swarm;

crates/ws/src/workers/socket.rs

@@ -0,0 +1,983 @@
+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 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<usize> = Default::default() }
+
+struct PendingScrapeResponse {
+    pending_worker_out_messages: usize,
+    stats: HashMap<InfoHash, ScrapeStatistics>,
+}
+
+struct ConnectionReference {
+    task_handle: Option<JoinHandle<()>>,
+    /// Sender part of channel used to pass on outgoing messages from request
+    /// worker
+    out_message_sender: Rc<LocalSender<(OutMessageMeta, OutMessage)>>,
+    /// Updated after sending message to peer
+    valid_until: ValidUntil,
+    announced_info_hashes: HashMap<InfoHash, PeerId>,
+    ip_version: IpVersion,
+    opt_peer_client: Option<(PeerClient, String)>,
+}
+
+pub async fn run_socket_worker(
+    _sentinel: PanicSentinel,
+    config: Config,
+    state: State,
+    opt_tls_config: Option<Arc<RustlsConfig>>,
+    control_message_mesh_builder: MeshBuilder<SwarmControlMessage, Partial>,
+    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) move || {
+            clean_connections(
+                config.clone(),
+                connection_slab.clone(),
+                server_start_instant,
+            )
+        }),
+        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,
+                });
+
+                ::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<Config>,
+    connection_slab: Rc<RefCell<Slab<ConnectionReference>>>,
+    server_start_instant: ServerStartInstant,
+) -> Option<Duration> {
+    let now = server_start_instant.seconds_elapsed();
+
+    connection_slab.borrow_mut().retain(|_, reference| {
+        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<RefCell<Slab<ConnectionReference>>>,
+) {
+    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<Config>,
+    access_list: Arc<AccessListArcSwap>,
+    in_message_senders: Rc<Senders<(InMessageMeta, InMessage)>>,
+    tq_prioritized: TaskQueueHandle,
+    tq_regular: TaskQueueHandle,
+    connection_slab: Rc<RefCell<Slab<ConnectionReference>>>,
+    out_message_sender: Rc<LocalSender<(OutMessageMeta, OutMessage)>>,
+    out_message_receiver: LocalReceiver<(OutMessageMeta, OutMessage)>,
+    server_start_instant: ServerStartInstant,
+    out_message_consumer_id: ConsumerId,
+    connection_id: ConnectionId,
+    opt_tls_config: Option<Arc<RustlsConfig>>,
+    ip_version: IpVersion,
+    mut stream: TcpStream,
+) -> anyhow::Result<()> {
+    if let Some(tls_config) = opt_tls_config {
+        let tls_acceptor: TlsAcceptor = tls_config.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<Config>,
+    access_list: Arc<AccessListArcSwap>,
+    in_message_senders: Rc<Senders<(InMessageMeta, InMessage)>>,
+    tq_prioritized: TaskQueueHandle,
+    tq_regular: TaskQueueHandle,
+    connection_slab: Rc<RefCell<Slab<ConnectionReference>>>,
+    out_message_sender: Rc<LocalSender<(OutMessageMeta, OutMessage)>>,
+    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<S> {
+    config: Rc<Config>,
+    access_list_cache: AccessListCache,
+    connection_slab: Rc<RefCell<Slab<ConnectionReference>>>,
+    in_message_senders: Rc<Senders<(InMessageMeta, InMessage)>>,
+    out_message_sender: Rc<LocalSender<(OutMessageMeta, OutMessage)>>,
+    pending_scrape_slab: Rc<RefCell<Slab<PendingScrapeResponse>>>,
+    out_message_consumer_id: ConsumerId,
+    ws_in: SplitStream<WebSocketStream<S>>,
+    ip_version: IpVersion,
+    connection_id: ConnectionId,
+}
+
+impl<S: futures::AsyncRead + futures::AsyncWrite + Unpin> ConnectionReader<S> {
+    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<usize, Vec<InfoHash>> = 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,
+                        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<ErrorResponseAction>,
+        info_hash: Option<InfoHash>,
+    ) -> 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<PendingScrapeId>) -> 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<S> {
+    config: Rc<Config>,
+    out_message_receiver: LocalReceiver<(OutMessageMeta, OutMessage)>,
+    connection_slab: Rc<RefCell<Slab<ConnectionReference>>>,
+    ws_out: SplitSink<WebSocketStream<S>, tungstenite::Message>,
+    pending_scrape_slab: Rc<RefCell<Slab<PendingScrapeResponse>>>,
+    server_start_instant: ServerStartInstant,
+    connection_id: ConnectionId,
+    ip_version: IpVersion,
+}
+
+impl<S: futures::AsyncRead + futures::AsyncWrite + Unpin> ConnectionWriter<S> {
+    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,
+                                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::Offer(_) => "offer",
+                        OutMessage::Answer(_) => "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<TcpListener> {
+    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",
+    }
+}

crates/ws/src/workers/swarm.rs

@@ -0,0 +1,539 @@
+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<usize> = 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<usize>) -> 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<PeerId, Peer>;
+
+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<InfoHash, TorrentData>;
+
+#[derive(Default)]
+struct TorrentMaps {
+    pub ipv4: TorrentMap,
+    pub ipv6: TorrentMap,
+}
+
+impl TorrentMaps {
+    fn clean(
+        &mut self,
+        config: &Config,
+        access_list: &Arc<AccessListArcSwap>,
+        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<SwarmControlMessage, Partial>,
+    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<S>(torrents: Rc<RefCell<TorrentMaps>>, mut stream: S)
+where
+    S: futures_lite::Stream<Item = SwarmControlMessage> + ::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<S>(
+    config: Config,
+    torrents: Rc<RefCell<TorrentMaps>>,
+    server_start_instant: ServerStartInstant,
+    out_message_senders: Rc<Senders<(OutMessageMeta, OutMessage)>>,
+    stream: S,
+) where
+    S: futures_lite::Stream<Item = (InMessageMeta, InMessage)> + ::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<Peer> = 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 middleman_offer = MiddlemanOfferToPeer {
+                action: AnnounceAction,
+                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::Offer(middleman_offer)));
+            ::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.to_peer_id, request.offer_id)
+    {
+        if let Some(answer_receiver) = torrent_data.peers.get(&answer_receiver_id) {
+            let middleman_answer = MiddlemanAnswerToPeer {
+                action: AnnounceAction,
+                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::Answer(middleman_answer)));
+            ::log::trace!("sending middleman answer to {:?}", meta);
+        }
+    }
+
+    let out_message = OutMessage::AnnounceResponse(AnnounceResponse {
+        action: AnnounceAction,
+        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,
+        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)));
+}