diff --git a/Cargo.lock b/Cargo.lock index 24bc347..767d617 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2032,6 +2032,7 @@ dependencies = [ "anyhow", "clap", "image", + "libc", "obscura-browser", "obscura-cdp", "obscura-dom", diff --git a/crates/obscura-cli/Cargo.toml b/crates/obscura-cli/Cargo.toml index dc99798..bebc740 100644 --- a/crates/obscura-cli/Cargo.toml +++ b/crates/obscura-cli/Cargo.toml @@ -34,5 +34,8 @@ url = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } +[target.'cfg(unix)'.dependencies] +libc = "0.2" + [dev-dependencies] image = { version = "0.25", default-features = false, features = ["png"] } diff --git a/crates/obscura-cli/src/main.rs b/crates/obscura-cli/src/main.rs index fd19f24..e0522bc 100644 --- a/crates/obscura-cli/src/main.rs +++ b/crates/obscura-cli/src/main.rs @@ -83,7 +83,8 @@ enum Command { /// Maximum live CDP connections. Each connection runs on its own OS /// thread with its own V8 isolates, so this bounds the server's thread /// and memory footprint. Connections beyond the limit are refused with - /// a 503 rather than queued. + /// a 503 rather than queued. With --workers N the limit applies to each + /// worker process, so the server-wide maximum is N times this value. #[arg(long, default_value_t = obscura_cdp::DEFAULT_MAX_CONNECTIONS)] max_connections: usize, @@ -473,6 +474,7 @@ async fn run_cli() -> anyhow::Result<()> { stealth, user_agent, font_dirs, + max_connections, ) .await?; } else { @@ -637,12 +639,17 @@ async fn run_multi_worker_serve( stealth: bool, user_agent: Option, font_dirs: Vec, + max_connections: usize, ) -> anyhow::Result<()> { use tokio::io::AsyncWriteExt as _; use tokio::net::TcpListener; use std::sync::atomic::{AtomicBool, Ordering}; let exe = std::env::current_exe()?; + // Register before spawning workers so an early SIGTERM is not fatal to the + // balancer alone and leaves no orphans. + let shutdown = shutdown_signal()?; + tokio::pin!(shutdown); // Claim the public port before starting children so another process cannot // take it during worker startup. let listener = TcpListener::bind((host.as_str(), port)).await?; @@ -663,6 +670,9 @@ async fn run_multi_worker_serve( let mut cmd = TokioCommand::new(&exe); cmd.kill_on_drop(true); cmd.arg("serve").arg("--port").arg(worker_port.to_string()); + // Same limit per worker, not split: connections are not spread evenly, + // so a split would refuse clients while other workers have room. + cmd.arg("--max-connections").arg(max_connections.to_string()); // Workers receive the client-facing Host header through the TCP // load balancer. Let their CDP security gate accept that public port // while it continues to reject foreign hosts and browser origins. @@ -703,18 +713,29 @@ async fn run_multi_worker_serve( let mut availability = Vec::with_capacity(workers as usize); let mut supervisors = tokio::task::JoinSet::new(); + let (stop_tx, stop_rx) = tokio::sync::watch::channel(()); for ((mut child, mut cmd), &worker_port) in children.into_iter().zip(&worker_ports) { let ready = Arc::new(AtomicBool::new(true)); availability.push(ready.clone()); + let mut stop = stop_rx.clone(); supervisors.spawn(async move { loop { - let status = child.wait().await; + let status = tokio::select! { + status = child.wait() => status, + _ = stop.changed() => { + stop_worker(&mut child).await; + return; + } + }; ready.store(false, Ordering::Relaxed); tracing::warn!("worker on port {} exited: {:?}", worker_port, status); loop { // ponytail: fixed retry bounds crash loops; add backoff if // persistently failing worker configurations need it. tokio::time::sleep(Duration::from_millis(100)).await; + if stop.has_changed().unwrap_or(true) { + return; + } match cmd.spawn() { Ok(mut replacement) => { let deadline = tokio::time::Instant::now() + Duration::from_secs(5); @@ -744,7 +765,10 @@ async fn run_multi_worker_serve( let mut next_worker = 0usize; loop { - let (client_stream, peer_addr) = listener.accept().await?; + let (client_stream, peer_addr) = tokio::select! { + accepted = listener.accept() => accepted?, + _ = &mut shutdown => break, + }; if let Err(error) = client_stream.set_nodelay(true) { tracing::warn!("client {} TCP_NODELAY failed: {}", peer_addr, error); } @@ -858,6 +882,56 @@ async fn run_multi_worker_serve( } }); } + + tracing::info!("Shutting down, stopping workers"); + drop(stop_tx); + // Workers are stopped in parallel by their supervisors; anything left after + // the grace period is aborted, and kill_on_drop then kills the child. + let _ = tokio::time::timeout(Duration::from_secs(5), async { + while supervisors.join_next().await.is_some() {} + }) + .await; + Ok(()) +} + +/// Resolves on SIGTERM or SIGINT (Ctrl-C elsewhere). Handlers are installed +/// when this is called, not on first poll. +fn shutdown_signal() -> std::io::Result> { + #[cfg(unix)] + { + use tokio::signal::unix::{signal, SignalKind}; + let mut term = signal(SignalKind::terminate())?; + let mut int = signal(SignalKind::interrupt())?; + Ok(async move { + tokio::select! { + _ = term.recv() => {} + _ = int.recv() => {} + } + }) + } + #[cfg(not(unix))] + { + Ok(async { + let _ = tokio::signal::ctrl_c().await; + }) + } +} + +/// Ask a worker to exit (SIGTERM lets it flush cookies), then kill it if it is +/// still running after a short grace period. +async fn stop_worker(child: &mut tokio::process::Child) { + #[cfg(unix)] + if let Some(pid) = child.id() { + // SAFETY: plain signal send to a child we still own (not yet reaped). + unsafe { libc::kill(pid as libc::pid_t, libc::SIGTERM) }; + if tokio::time::timeout(Duration::from_secs(3), child.wait()) + .await + .is_ok() + { + return; + } + } + let _ = child.kill().await; } async fn settle_page(page: &mut Page, wait_secs: u64, fixed: bool) { diff --git a/crates/obscura-cli/tests/serve_shutdown.rs b/crates/obscura-cli/tests/serve_shutdown.rs new file mode 100644 index 0000000..60a1a00 --- /dev/null +++ b/crates/obscura-cli/tests/serve_shutdown.rs @@ -0,0 +1,79 @@ +#![cfg(target_os = "linux")] + +use std::io::{Read, Write}; +use std::net::{TcpListener, TcpStream}; +use std::os::unix::process::CommandExt; +use std::path::Path; +use std::process::{Child, Command, Stdio}; +use std::time::{Duration, Instant}; + +struct Server(Child); + +impl Drop for Server { + fn drop(&mut self) { + let _ = Command::new("kill").args(["-KILL", "--", &format!("-{}", self.0.id())]) + .stdout(Stdio::null()).stderr(Stdio::null()).status(); + let _ = self.0.wait(); + } +} + +fn workers(parent: u32) -> Vec { + let root = format!("/proc/{parent}/task"); + let mut children = Vec::new(); + for task in std::fs::read_dir(root).unwrap() { + if let Ok(ids) = std::fs::read_to_string(task.unwrap().path().join("children")) { + children.extend(ids.split_whitespace().map(|id| id.parse::().unwrap())); + } + } + children.sort_unstable(); + children.dedup(); + children +} + +fn discovery_ok(port: u16) -> bool { + let Ok(mut stream) = TcpStream::connect(("127.0.0.1", port)) else { return false }; + stream.set_read_timeout(Some(Duration::from_secs(2))).unwrap(); + write!(stream, "GET /json/version HTTP/1.1\r\nHost: 127.0.0.1:{port}\r\nConnection: close\r\n\r\n").unwrap(); + let mut response = String::new(); + stream.read_to_string(&mut response).is_ok() && response.starts_with("HTTP/1.1 200") +} + +#[test] +fn balancer_signal_stops_workers_and_workers_get_max_connections() { + for signal in ["-TERM", "-INT"] { + let reservation = TcpListener::bind(("127.0.0.1", 0)).unwrap(); + let port = reservation.local_addr().unwrap().port(); + drop(reservation); + let mut server = Server(Command::new(env!("CARGO_BIN_EXE_obscura")) + .args(["serve", "--host", "127.0.0.1", "--port", &port.to_string(), + "--workers", "2", "--max-connections", "7"]) + .process_group(0).stdout(Stdio::null()).stderr(Stdio::null()).spawn().unwrap()); + let deadline = Instant::now() + Duration::from_secs(10); + while !discovery_ok(port) { + assert!(server.0.try_wait().unwrap().is_none(), "balancer exited"); + assert!(Instant::now() < deadline, "balancer did not become ready"); + std::thread::sleep(Duration::from_millis(10)); + } + let pids = workers(server.0.id()); + assert_eq!(pids.len(), 2); + for pid in &pids { + let argv = std::fs::read(format!("/proc/{pid}/cmdline")).unwrap(); + let argv: Vec<&[u8]> = argv.split(|b| *b == 0).collect(); + assert!(argv.windows(2).any(|w| w == [&b"--max-connections"[..], b"7"]), + "worker {pid} missing --max-connections 7: {argv:?}"); + } + + // Signal only the balancer, not the process group. + assert!(Command::new("kill").args([signal, &server.0.id().to_string()]).status().unwrap().success()); + let deadline = Instant::now() + Duration::from_secs(10); + while server.0.try_wait().unwrap().is_none() { + assert!(Instant::now() < deadline, "balancer ignored {signal}"); + std::thread::sleep(Duration::from_millis(10)); + } + let deadline = Instant::now() + Duration::from_secs(10); + while pids.iter().any(|pid| Path::new(&format!("/proc/{pid}")).exists()) { + assert!(Instant::now() < deadline, "workers survived {signal}: {pids:?}"); + std::thread::sleep(Duration::from_millis(10)); + } + } +}