diff --git a/aquatic_http/src/lib.rs b/aquatic_http/src/lib.rs index 7ce2065..f1b8ca7 100644 --- a/aquatic_http/src/lib.rs +++ b/aquatic_http/src/lib.rs @@ -6,10 +6,14 @@ use aquatic_common::{ }, privileges::PrivilegeDropper, rustls_config::create_rustls_config, + PanicSentinelWatcher, }; use common::State; use glommio::{channels::channel_mesh::MeshBuilder, prelude::*}; -use signal_hook::{consts::SIGUSR1, iterator::Signals}; +use signal_hook::{ + consts::{SIGTERM, SIGUSR1}, + iterator::Signals, +}; use std::sync::Arc; use crate::config::Config; @@ -24,45 +28,18 @@ pub const APP_VERSION: &str = env!("CARGO_PKG_VERSION"); const SHARED_CHANNEL_SIZE: usize = 1024; pub fn run(config: Config) -> ::anyhow::Result<()> { + let mut signals = Signals::new([SIGUSR1, SIGTERM])?; + let state = State::default(); update_access_list(&config.access_list, &state.access_list)?; - let mut signals = Signals::new(::std::iter::once(SIGUSR1))?; - - { - let config = config.clone(); - let state = state.clone(); - - ::std::thread::spawn(move || run_inner(config, state)); - } - - if config.cpu_pinning.active { - set_affinity_for_util_worker( - &config.cpu_pinning, - config.socket_workers, - config.request_workers, - )?; - } - - for signal in &mut signals { - match signal { - SIGUSR1 => { - let _ = update_access_list(&config.access_list, &state.access_list); - } - _ => unreachable!(), - } - } - - Ok(()) -} - -pub fn run_inner(config: Config, state: State) -> anyhow::Result<()> { let num_peers = config.socket_workers + config.request_workers; let request_mesh_builder = MeshBuilder::partial(num_peers, SHARED_CHANNEL_SIZE); let response_mesh_builder = MeshBuilder::partial(num_peers, SHARED_CHANNEL_SIZE); + let (sentinel_watcher, sentinel) = PanicSentinelWatcher::create_with_sentinel(); let priv_dropper = PrivilegeDropper::new(config.privileges.clone(), config.socket_workers); let tls_config = Arc::new(create_rustls_config( @@ -73,6 +50,7 @@ pub fn run_inner(config: Config, state: State) -> anyhow::Result<()> { 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 tls_config = tls_config.clone(); @@ -88,22 +66,26 @@ pub fn run_inner(config: Config, state: State) -> anyhow::Result<()> { )?; let builder = LocalExecutorBuilder::new(placement).name("socket"); - let executor = builder.spawn(move || async move { - workers::socket::run_socket_worker( - config, - state, - tls_config, - request_mesh_builder, - response_mesh_builder, - priv_dropper, - ) - .await - }); + let executor = builder + .spawn(move || async move { + workers::socket::run_socket_worker( + sentinel, + config, + state, + tls_config, + request_mesh_builder, + response_mesh_builder, + priv_dropper, + ) + .await + }) + .map_err(|err| anyhow::anyhow!("Spawning executor failed: {:#}", err))?; executors.push(executor); } for i in 0..(config.request_workers) { + let sentinel = sentinel.clone(); let config = config.clone(); let state = state.clone(); let request_mesh_builder = request_mesh_builder.clone(); @@ -117,15 +99,18 @@ pub fn run_inner(config: Config, state: State) -> anyhow::Result<()> { )?; let builder = LocalExecutorBuilder::new(placement).name("request"); - let executor = builder.spawn(move || async move { - workers::request::run_request_worker( - config, - state, - request_mesh_builder, - response_mesh_builder, - ) - .await - }); + let executor = builder + .spawn(move || async move { + workers::request::run_request_worker( + sentinel, + config, + state, + request_mesh_builder, + response_mesh_builder, + ) + .await + }) + .map_err(|err| anyhow::anyhow!("Spawning executor failed: {:#}", err))?; executors.push(executor); } @@ -138,11 +123,20 @@ pub fn run_inner(config: Config, state: State) -> anyhow::Result<()> { )?; } - for executor in executors { - executor - .expect("failed to spawn local executor") - .join() - .unwrap(); + 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(()) diff --git a/aquatic_http/src/workers/request.rs b/aquatic_http/src/workers/request.rs index 8ac677e..e664c4e 100644 --- a/aquatic_http/src/workers/request.rs +++ b/aquatic_http/src/workers/request.rs @@ -17,8 +17,8 @@ use rand::SeedableRng; use smartstring::{LazyCompact, SmartString}; use aquatic_common::access_list::{create_access_list_cache, AccessListArcSwap, AccessListCache}; -use aquatic_common::extract_response_peers; use aquatic_common::ValidUntil; +use aquatic_common::{extract_response_peers, PanicSentinel}; use aquatic_common::{AmortizedIndexMap, CanonicalSocketAddr}; use aquatic_http_protocol::common::*; use aquatic_http_protocol::request::*; @@ -175,6 +175,7 @@ impl TorrentMaps { } pub async fn run_request_worker( + _sentinel: PanicSentinel, config: Config, state: State, request_mesh_builder: MeshBuilder, diff --git a/aquatic_http/src/workers/socket.rs b/aquatic_http/src/workers/socket.rs index 5271ec0..4b7a7f6 100644 --- a/aquatic_http/src/workers/socket.rs +++ b/aquatic_http/src/workers/socket.rs @@ -9,7 +9,7 @@ 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::CanonicalSocketAddr; +use aquatic_common::{CanonicalSocketAddr, PanicSentinel}; use aquatic_http_protocol::common::InfoHash; use aquatic_http_protocol::request::{Request, RequestParseError, ScrapeRequest}; use aquatic_http_protocol::response::{ @@ -54,6 +54,7 @@ struct ConnectionReference { } pub async fn run_socket_worker( + _sentinel: PanicSentinel, config: Config, state: State, tls_config: Arc,