use std::io;
use std::net::{TcpListener, TcpStream};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicI32, AtomicU8, Ordering};
use std::time::{Duration, Instant};
use super::config::{LogLevel, ServerConfig};
use super::handler::TransformOptionsPayload;
use super::routing::handle_stream;
use super::stderr_write;
pub(super) const SOCKET_READ_TIMEOUT: Duration = Duration::from_secs(60);
pub(super) const SOCKET_WRITE_TIMEOUT: Duration = Duration::from_secs(60);
pub(super) const HEADER_READ_DEADLINE: Duration = Duration::from_secs(15);
const WORKER_THREADS: usize = 8;
fn worker_pool_size(max_concurrent_transforms: u64) -> usize {
usize::try_from(max_concurrent_transforms)
.unwrap_or(usize::MAX)
.saturating_add(WORKER_THREADS)
}
pub fn serve(listener: TcpListener) -> io::Result<()> {
let config = ServerConfig::from_env()?;
for (ok, name) in super::handler::storage_health_check(&config) {
if !ok {
return Err(io::Error::new(
io::ErrorKind::ConnectionRefused,
format!(
"storage connectivity check failed for `{name}` — verify the backend \
endpoint, credentials, and container/bucket configuration"
),
));
}
}
serve_with_config(listener, config)
}
pub fn serve_with_config(listener: TcpListener, config: ServerConfig) -> io::Result<()> {
let config = Arc::new(config);
let (sender, receiver) = std::sync::mpsc::channel::<(TcpStream, Instant)>();
let receiver = Arc::new(std::sync::Mutex::new(receiver));
let pool_size = worker_pool_size(config.max_concurrent_transforms);
let mut workers = Vec::with_capacity(pool_size);
for _ in 0..pool_size {
let rx = Arc::clone(&receiver);
let cfg = Arc::clone(&config);
workers.push(std::thread::spawn(move || {
loop {
let (stream, accepted_at) = {
let guard = rx.lock().expect("worker lock poisoned");
match guard.recv() {
Ok(accepted) => accepted,
Err(_) => break,
}
}; if let Err(err) = handle_stream(stream, accepted_at, &cfg) {
cfg.log_warn(&format!("failed to handle connection: {err}"));
}
}
}));
}
let (shutdown_read_fd, shutdown_write_fd) = create_shutdown_pipe()?;
install_signal_handler(
Arc::clone(&config.draining),
shutdown_write_fd,
Arc::clone(&config.log_level),
);
if let Some(ref path) = config.presets_file_path {
let presets = Arc::clone(&config.presets);
let draining = Arc::clone(&config.draining);
let cfg = Arc::clone(&config);
let path = path.clone();
std::thread::Builder::new()
.name("preset-watcher".into())
.spawn(move || preset_watcher(presets, path, draining, cfg))
.expect("failed to spawn preset watcher thread");
}
listener.set_nonblocking(true)?;
let mut drain_deadline: Option<Instant> = None;
loop {
let remaining = drain_deadline
.map(|deadline: Instant| deadline.saturating_duration_since(Instant::now()));
if remaining.is_some_and(|left| left.is_zero()) {
break;
}
wait_for_accept_or_shutdown(&listener, shutdown_read_fd, &config.draining, remaining);
let signalled =
poll_shutdown_pipe(shutdown_read_fd) || config.draining.load(Ordering::SeqCst);
if signalled && drain_deadline.is_none() {
let drain_secs = config.shutdown_drain_secs;
config.log(&format!(
"shutdown: drain started, waiting {drain_secs}s for load balancers"
));
if drain_secs == 0 {
break;
}
drain_deadline = Some(Instant::now() + Duration::from_secs(drain_secs));
}
match listener.accept() {
Ok((stream, _addr)) => {
let _ = stream.set_nonblocking(false);
if sender.send((stream, Instant::now())).is_err() {
break;
}
}
Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
}
Err(err) => return Err(err),
}
}
config.log("shutdown: drain complete, closing listener");
drop(listener);
drop(sender);
let deadline = Instant::now() + Duration::from_secs(15);
for worker in workers {
let remaining = deadline.saturating_duration_since(Instant::now());
if remaining.is_zero() {
stderr_write("shutdown: timed out waiting for worker threads");
break;
}
let worker_done =
std::sync::Arc::new((std::sync::Mutex::new(false), std::sync::Condvar::new()));
let wd = std::sync::Arc::clone(&worker_done);
std::thread::spawn(move || {
let _ = worker.join();
let (lock, cvar) = &*wd;
*lock.lock().expect("shutdown notify lock") = true;
cvar.notify_one();
});
let (lock, cvar) = &*worker_done;
let mut done = lock.lock().expect("shutdown wait lock");
while !*done {
let (guard, timeout) = cvar
.wait_timeout(done, remaining)
.expect("shutdown condvar wait");
done = guard;
if timeout.timed_out() {
stderr_write("shutdown: timed out waiting for a worker thread");
break;
}
}
}
config.log("shutdown: complete");
close_shutdown_pipe(shutdown_read_fd, shutdown_write_fd);
Ok(())
}
pub fn serve_once(listener: TcpListener) -> io::Result<()> {
let config = ServerConfig::from_env()?;
serve_once_with_config(listener, config)
}
pub fn serve_once_with_config(listener: TcpListener, config: ServerConfig) -> io::Result<()> {
let (stream, _) = listener.accept()?;
let accepted_at = Instant::now();
handle_stream(stream, accepted_at, &config)
}
#[cfg(unix)]
fn create_shutdown_pipe() -> io::Result<(i32, i32)> {
let mut fds = [0i32; 2];
if unsafe { libc::pipe(fds.as_mut_ptr()) } != 0 {
return Err(io::Error::last_os_error());
}
unsafe {
libc::fcntl(fds[0], libc::F_SETFL, libc::O_NONBLOCK);
libc::fcntl(fds[1], libc::F_SETFL, libc::O_NONBLOCK);
}
Ok((fds[0], fds[1]))
}
#[cfg(windows)]
fn create_shutdown_pipe() -> io::Result<(i32, i32)> {
Ok((-1, -1))
}
#[cfg(unix)]
fn poll_shutdown_pipe(read_fd: i32) -> bool {
let mut buf = [0u8; 1];
let n = unsafe { libc::read(read_fd, buf.as_mut_ptr().cast(), 1) };
n > 0
}
#[cfg(windows)]
fn poll_shutdown_pipe(_read_fd: i32) -> bool {
false
}
#[cfg(unix)]
fn wait_for_accept_or_shutdown(
listener: &std::net::TcpListener,
shutdown_read_fd: i32,
_draining: &AtomicBool,
timeout: Option<Duration>,
) {
use std::os::unix::io::AsRawFd;
let listener_fd = listener.as_raw_fd();
let mut fds = [
libc::pollfd {
fd: listener_fd,
events: libc::POLLIN,
revents: 0,
},
libc::pollfd {
fd: shutdown_read_fd,
events: libc::POLLIN,
revents: 0,
},
];
let timeout_ms = match timeout {
None => -1,
Some(remaining) => i32::try_from(remaining.as_millis()).unwrap_or(i32::MAX),
};
unsafe { libc::poll(fds.as_mut_ptr(), 2, timeout_ms) };
}
#[cfg(windows)]
fn wait_for_accept_or_shutdown(
_listener: &std::net::TcpListener,
_shutdown_read_fd: i32,
draining: &AtomicBool,
timeout: Option<Duration>,
) {
const NAP: Duration = Duration::from_millis(10);
match timeout {
Some(remaining) if remaining.is_zero() => {}
Some(remaining) => std::thread::sleep(NAP.min(remaining)),
None => {
if !draining.load(Ordering::SeqCst) {
std::thread::sleep(NAP);
}
}
}
}
#[cfg(unix)]
fn close_shutdown_pipe(read_fd: i32, write_fd: i32) {
unsafe {
libc::close(read_fd);
libc::close(write_fd);
}
}
#[cfg(windows)]
fn close_shutdown_pipe(_read_fd: i32, _write_fd: i32) {}
static SHUTDOWN_PIPE_WR: AtomicI32 = AtomicI32::new(-1);
static GLOBAL_DRAINING: std::sync::atomic::AtomicPtr<AtomicBool> =
std::sync::atomic::AtomicPtr::new(std::ptr::null_mut());
static GLOBAL_LOG_LEVEL: std::sync::atomic::AtomicPtr<AtomicU8> =
std::sync::atomic::AtomicPtr::new(std::ptr::null_mut());
#[cfg(unix)]
fn install_signal_handler(draining: Arc<AtomicBool>, write_fd: i32, log_level: Arc<AtomicU8>) {
SHUTDOWN_PIPE_WR.store(write_fd, Ordering::SeqCst);
let ptr = Arc::into_raw(draining).cast_mut();
GLOBAL_DRAINING.store(ptr, Ordering::SeqCst);
let lvl_ptr = Arc::into_raw(log_level).cast_mut();
GLOBAL_LOG_LEVEL.store(lvl_ptr, Ordering::SeqCst);
unsafe {
let mut sa: libc::sigaction = std::mem::zeroed();
sa.sa_sigaction = signal_handler as *const () as libc::sighandler_t;
libc::sigemptyset(&mut sa.sa_mask);
sa.sa_flags = libc::SA_RESTART;
libc::sigaction(libc::SIGTERM, &sa, std::ptr::null_mut());
libc::sigaction(libc::SIGINT, &sa, std::ptr::null_mut());
let mut sa_usr1: libc::sigaction = std::mem::zeroed();
sa_usr1.sa_sigaction = sigusr1_handler as *const () as libc::sighandler_t;
libc::sigemptyset(&mut sa_usr1.sa_mask);
sa_usr1.sa_flags = libc::SA_RESTART;
libc::sigaction(libc::SIGUSR1, &sa_usr1, std::ptr::null_mut());
}
}
#[cfg(unix)]
extern "C" fn sigusr1_handler(_sig: libc::c_int) {
let ptr = GLOBAL_LOG_LEVEL.load(Ordering::SeqCst);
if ptr.is_null() {
return;
}
let level_atomic = unsafe { &*ptr };
let current = level_atomic.load(Ordering::SeqCst);
let next = LogLevel::from_u8(current).cycle();
level_atomic.store(next as u8, Ordering::SeqCst);
let msg = match next {
LogLevel::Error => b"[log] level changed to error\n" as &[u8],
LogLevel::Warn => b"[log] level changed to warn\n",
LogLevel::Info => b"[log] level changed to info\n",
LogLevel::Debug => b"[log] level changed to debug\n",
};
unsafe { libc::write(2, msg.as_ptr().cast(), msg.len()) };
}
#[cfg(unix)]
extern "C" fn signal_handler(_sig: libc::c_int) {
let ptr = GLOBAL_DRAINING.load(Ordering::SeqCst);
if !ptr.is_null() {
unsafe { (*ptr).store(true, Ordering::SeqCst) };
}
let fd = SHUTDOWN_PIPE_WR.load(Ordering::SeqCst);
if fd >= 0 {
let byte: u8 = 1;
unsafe { libc::write(fd, (&byte as *const u8).cast(), 1) };
}
}
#[cfg(windows)]
fn install_signal_handler(draining: Arc<AtomicBool>, _write_fd: i32, _log_level: Arc<AtomicU8>) {
let ptr = Arc::into_raw(draining).cast_mut();
GLOBAL_DRAINING.store(ptr, Ordering::SeqCst);
unsafe {
libc::signal(libc::SIGINT, windows_signal_handler as libc::sighandler_t);
}
}
#[cfg(windows)]
extern "C" fn windows_signal_handler(_sig: libc::c_int) {
let ptr = GLOBAL_DRAINING.load(Ordering::SeqCst);
if !ptr.is_null() {
unsafe { (*ptr).store(true, Ordering::SeqCst) };
}
unsafe {
libc::signal(libc::SIGINT, windows_signal_handler as libc::sighandler_t);
}
}
const PRESET_WATCH_INTERVAL: Duration = Duration::from_secs(5);
pub(super) fn preset_watcher(
presets: Arc<std::sync::RwLock<std::collections::HashMap<String, TransformOptionsPayload>>>,
path: std::path::PathBuf,
draining: Arc<AtomicBool>,
config: Arc<ServerConfig>,
) {
use super::config::parse_presets_file;
use std::fs;
let mut last_modified = fs::metadata(&path).and_then(|m| m.modified()).ok();
loop {
std::thread::sleep(PRESET_WATCH_INTERVAL);
if draining.load(Ordering::Relaxed) {
break;
}
let current_modified = match fs::metadata(&path).and_then(|m| m.modified()) {
Ok(mtime) => Some(mtime),
Err(err) => {
config.log_warn(&format!(
"[presets] failed to stat `{}`: {err}",
path.display()
));
continue;
}
};
if current_modified == last_modified {
continue;
}
match parse_presets_file(&path) {
Ok(new_presets) => {
let count = new_presets.len();
*presets.write().expect("presets lock poisoned") = new_presets;
last_modified = current_modified;
config.log(&format!(
"[presets] reloaded {count} presets from `{}`",
path.display()
));
}
Err(err) => {
config.log_warn(&format!(
"[presets] reload failed for `{}`: {err} (keeping previous presets)",
path.display()
));
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[cfg(unix)]
use serial_test::serial;
#[cfg(unix)]
use std::net::TcpListener;
#[cfg(unix)]
use std::sync::atomic::AtomicU8;
#[cfg(unix)]
struct ShutdownSignalGuard {
previous_draining: *mut AtomicBool,
previous_write_fd: i32,
draining: *mut AtomicBool,
read_fd: i32,
write_fd: i32,
}
#[cfg(unix)]
impl Drop for ShutdownSignalGuard {
fn drop(&mut self) {
GLOBAL_DRAINING.store(self.previous_draining, Ordering::SeqCst);
SHUTDOWN_PIPE_WR.store(self.previous_write_fd, Ordering::SeqCst);
close_shutdown_pipe(self.read_fd, self.write_fd);
unsafe { drop(Box::from_raw(self.draining)) };
}
}
#[cfg(unix)]
struct LogLevelGuard {
previous: *mut AtomicU8,
log_level: *mut AtomicU8,
}
#[cfg(unix)]
impl Drop for LogLevelGuard {
fn drop(&mut self) {
GLOBAL_LOG_LEVEL.store(self.previous, Ordering::SeqCst);
unsafe { drop(Box::from_raw(self.log_level)) };
}
}
#[test]
fn the_worker_pool_exceeds_the_transform_limit_at_every_limit() {
for limit in [1_u64, 2, 4, 7, 8, 9, 32, 1024] {
let pool = worker_pool_size(limit);
assert!(
pool >= limit as usize + WORKER_THREADS,
"limit {limit} sized the pool at {pool}"
);
}
}
#[cfg(unix)]
#[test]
fn wait_for_accept_or_shutdown_returns_when_pipe_is_ready() {
let listener = TcpListener::bind("127.0.0.1:0").expect("bind test listener");
let (read_fd, write_fd) = create_shutdown_pipe().expect("create shutdown pipe");
let byte: u8 = 1;
let written = unsafe { libc::write(write_fd, (&byte as *const u8).cast(), 1) };
assert_eq!(written, 1, "write shutdown wakeup byte");
let draining = AtomicBool::new(false);
let start = Instant::now();
wait_for_accept_or_shutdown(&listener, read_fd, &draining, None);
assert!(
start.elapsed() < Duration::from_millis(200),
"wait_for_accept_or_shutdown should return immediately when the pipe is readable"
);
assert!(poll_shutdown_pipe(read_fd));
close_shutdown_pipe(read_fd, write_fd);
}
#[cfg(unix)]
#[test]
fn wait_for_accept_or_shutdown_honours_the_drain_timeout() {
let listener = TcpListener::bind("127.0.0.1:0").expect("bind test listener");
let (read_fd, write_fd) = create_shutdown_pipe().expect("create shutdown pipe");
let draining = AtomicBool::new(true);
let start = Instant::now();
wait_for_accept_or_shutdown(
&listener,
read_fd,
&draining,
Some(Duration::from_millis(50)),
);
assert!(
start.elapsed() < Duration::from_secs(2),
"a bounded wait must return without traffic on the listener"
);
close_shutdown_pipe(read_fd, write_fd);
}
#[cfg(unix)]
#[test]
#[serial]
fn signal_handler_sets_draining_and_wakes_shutdown_pipe() {
let (read_fd, write_fd) = create_shutdown_pipe().expect("create shutdown pipe");
let draining = Box::into_raw(Box::new(AtomicBool::new(false)));
let previous_draining = GLOBAL_DRAINING.swap(draining, Ordering::SeqCst);
let previous_write_fd = SHUTDOWN_PIPE_WR.swap(write_fd, Ordering::SeqCst);
let _restore = ShutdownSignalGuard {
previous_draining,
previous_write_fd,
draining,
read_fd,
write_fd,
};
signal_handler(libc::SIGTERM);
assert!(unsafe { &*draining }.load(Ordering::SeqCst));
assert!(poll_shutdown_pipe(read_fd));
}
#[cfg(unix)]
#[test]
#[serial]
fn sigusr1_handler_cycles_global_log_level() {
let log_level = Box::into_raw(Box::new(AtomicU8::new(LogLevel::Info as u8)));
let previous = GLOBAL_LOG_LEVEL.swap(log_level, Ordering::SeqCst);
let _restore = LogLevelGuard {
previous,
log_level,
};
sigusr1_handler(libc::SIGUSR1);
assert_eq!(
unsafe { &*log_level }.load(Ordering::SeqCst),
LogLevel::Debug as u8
);
}
}