diff --git a/aquatic_udp/src/lib.rs b/aquatic_udp/src/lib.rs index 5a805d8..663eba0 100644 --- a/aquatic_udp/src/lib.rs +++ b/aquatic_udp/src/lib.rs @@ -2,6 +2,7 @@ pub mod common; pub mod config; pub mod workers; +use aquatic_common::PanicSentinel; use config::Config; use std::collections::BTreeMap; @@ -14,7 +15,7 @@ use aquatic_common::privileges::PrivilegeDropper; use crossbeam_channel::{bounded, unbounded}; use aquatic_common::access_list::update_access_list; -use signal_hook::consts::SIGUSR1; +use signal_hook::consts::{SIGTERM, SIGUSR1}; use signal_hook::iterator::Signals; use common::{ConnectedRequestSender, ConnectedResponseSender, SocketWorkerIndex, State}; @@ -29,7 +30,7 @@ pub fn run(config: Config) -> ::anyhow::Result<()> { update_access_list(&config.access_list, &state.access_list)?; - let mut signals = Signals::new(::std::iter::once(SIGUSR1))?; + let mut signals = Signals::new([SIGUSR1, SIGTERM])?; let priv_dropper = PrivilegeDropper::new(config.privileges.clone(), config.socket_workers); @@ -79,6 +80,7 @@ pub fn run(config: Config) -> ::anyhow::Result<()> { ); workers::request::run_request_worker( + PanicSentinel, config, state, request_receiver, @@ -109,6 +111,7 @@ pub fn run(config: Config) -> ::anyhow::Result<()> { ); workers::socket::run_socket_worker( + PanicSentinel, state, config, i, @@ -135,7 +138,7 @@ pub fn run(config: Config) -> ::anyhow::Result<()> { WorkerIndex::Util, ); - workers::statistics::run_statistics_worker(config, state); + workers::statistics::run_statistics_worker(PanicSentinel, config, state); }) .with_context(|| "spawn statistics worker")?; } @@ -153,6 +156,9 @@ pub fn run(config: Config) -> ::anyhow::Result<()> { SIGUSR1 => { let _ = update_access_list(&config.access_list, &state.access_list); } + SIGTERM => { + break; + } _ => unreachable!(), } } diff --git a/aquatic_udp/src/workers/request.rs b/aquatic_udp/src/workers/request.rs index 447fab4..12eb6b9 100644 --- a/aquatic_udp/src/workers/request.rs +++ b/aquatic_udp/src/workers/request.rs @@ -11,6 +11,7 @@ use aquatic_common::access_list::create_access_list_cache; use aquatic_common::access_list::AccessListArcSwap; use aquatic_common::AmortizedIndexMap; use aquatic_common::CanonicalSocketAddr; +use aquatic_common::PanicSentinel; use aquatic_common::ValidUntil; use crossbeam_channel::Receiver; use rand::{rngs::SmallRng, SeedableRng}; @@ -121,6 +122,7 @@ impl TorrentMaps { } pub fn run_request_worker( + _sentinel: PanicSentinel, config: Config, state: State, request_receiver: Receiver<(SocketWorkerIndex, ConnectedRequest, CanonicalSocketAddr)>, diff --git a/aquatic_udp/src/workers/socket.rs b/aquatic_udp/src/workers/socket.rs index ac9c556..c8b5e05 100644 --- a/aquatic_udp/src/workers/socket.rs +++ b/aquatic_udp/src/workers/socket.rs @@ -14,8 +14,8 @@ use slab::Slab; use aquatic_common::access_list::create_access_list_cache; use aquatic_common::access_list::AccessListCache; -use aquatic_common::ValidUntil; use aquatic_common::{AmortizedIndexMap, CanonicalSocketAddr}; +use aquatic_common::{PanicSentinel, ValidUntil}; use aquatic_udp_protocol::*; use socket2::{Domain, Protocol, Socket, Type}; @@ -151,6 +151,7 @@ impl PendingScrapeResponseSlab { } pub fn run_socket_worker( + _sentinel: PanicSentinel, state: State, config: Config, token_num: usize, @@ -161,7 +162,8 @@ pub fn run_socket_worker( let mut rng = StdRng::from_entropy(); let mut buffer = [0u8; MAX_PACKET_SIZE]; - let mut socket = UdpSocket::from_std(create_socket(&config, priv_dropper).expect("create socket")); + let mut socket = + UdpSocket::from_std(create_socket(&config, priv_dropper).expect("create socket")); let mut poll = Poll::new().expect("create poll"); let interests = Interest::READABLE; @@ -517,7 +519,10 @@ fn send_response( } } -pub fn create_socket(config: &Config, priv_dropper: PrivilegeDropper) -> anyhow::Result<::std::net::UdpSocket> { +pub fn create_socket( + config: &Config, + priv_dropper: PrivilegeDropper, +) -> anyhow::Result<::std::net::UdpSocket> { let socket = if config.network.address.is_ipv4() { Socket::new(Domain::IPV4, Type::DGRAM, Some(Protocol::UDP))? } else { @@ -525,10 +530,14 @@ pub fn create_socket(config: &Config, priv_dropper: PrivilegeDropper) -> anyhow: }; if config.network.only_ipv6 { - socket.set_only_v6(true).with_context(|| "socket: set only ipv6")?; + socket + .set_only_v6(true) + .with_context(|| "socket: set only ipv6")?; } - socket.set_reuse_port(true).with_context(|| "socket: set reuse port")?; + socket + .set_reuse_port(true) + .with_context(|| "socket: set reuse port")?; socket .set_nonblocking(true) diff --git a/aquatic_udp/src/workers/statistics.rs b/aquatic_udp/src/workers/statistics.rs index 20ec0e7..b54cbe8 100644 --- a/aquatic_udp/src/workers/statistics.rs +++ b/aquatic_udp/src/workers/statistics.rs @@ -5,6 +5,7 @@ use std::sync::Arc; use std::time::{Duration, Instant}; use anyhow::Context; +use aquatic_common::PanicSentinel; use num_format::{Locale, ToFormattedString}; use serde::Serialize; use time::format_description::well_known::Rfc2822; @@ -135,7 +136,7 @@ struct TemplateData { peer_update_interval: String, } -pub fn run_statistics_worker(config: Config, state: State) { +pub fn run_statistics_worker(_sentinel: PanicSentinel, config: Config, state: State) { let tt = if config.statistics.write_html_to_file { let mut tt = TinyTemplate::new(); diff --git a/aquatic_udp_bench/src/main.rs b/aquatic_udp_bench/src/main.rs index 6f65bc9..046ab6d 100644 --- a/aquatic_udp_bench/src/main.rs +++ b/aquatic_udp_bench/src/main.rs @@ -7,6 +7,7 @@ //! Scrape: 1 873 545 requests/second, 533.75 ns/request //! ``` +use aquatic_common::PanicSentinel; use aquatic_udp::workers::request::run_request_worker; use crossbeam_channel::unbounded; use num_format::{Locale, ToFormattedString}; @@ -55,6 +56,7 @@ pub fn run(bench_config: BenchConfig) -> ::anyhow::Result<()> { ::std::thread::spawn(move || { run_request_worker( + PanicSentinel, config, state, request_receiver,