mirror of
https://github.com/YGGverse/aquatic.git
synced 2026-04-01 02:05:30 +00:00
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
This commit is contained in:
Joakim Frostegård
GitHub
parent
af9d5a55f6
commit
fe5ccf6646
19 changed files with 1770 additions and 1583 deletions
362
crates/ws/src/workers/socket/mod.rs
Normal file
362
crates/ws/src/workers/socket/mod.rs
Normal file
|
|
@ -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<ConnectionId, ConnectionHandle>;
|
||||
|
||||
const LOCAL_CHANNEL_SIZE: usize = 16;
|
||||
|
||||
#[cfg(feature = "metrics")]
|
||||
thread_local! { static WORKER_INDEX: ::std::cell::Cell<usize> = 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<LocalSender<(OutMessageMeta, OutMessage)>>,
|
||||
/// Updated after sending message to peer
|
||||
valid_until: Rc<RefCell<ValidUntil>>,
|
||||
/// The TLS config used for this connection
|
||||
opt_tls_config: Option<Arc<RustlsConfig>>,
|
||||
valid_until_after_tls_update: Option<ValidUntil>,
|
||||
}
|
||||
|
||||
pub async fn run_socket_worker(
|
||||
_sentinel: PanicSentinel,
|
||||
config: Config,
|
||||
state: State,
|
||||
opt_tls_config: Option<Arc<ArcSwap<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_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<Config>,
|
||||
connection_slab: Rc<RefCell<ConnectionHandles>>,
|
||||
server_start_instant: ServerStartInstant,
|
||||
opt_tls_config: Option<Arc<ArcSwap<RustlsConfig>>>,
|
||||
) -> Option<Duration> {
|
||||
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<RefCell<ConnectionHandles>>,
|
||||
) {
|
||||
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<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",
|
||||
}
|
||||
}
|
||||
|
||||
fn calculate_in_message_consumer_index(config: &Config, info_hash: InfoHash) -> usize {
|
||||
(info_hash.0[0] as usize) % config.swarm_workers
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue