swadhin / justapi public

p4: multi-worker prefork supervisor (justapi serve --workers N) (ADR-061)

swadhinbiswas committed 3 months ago

7a47aad

1 parent cf681c4

commit 7a47aad

8 files changed

+590 -135

Files

crates

justapi-cli

src

main.rs +264 -135 workers.rs +254 Cargo.toml +1

justapi-core

src

server

mod.rs +1 BENCHMARKS.md +19 Cargo.lock +1 Cargo.toml +1 DECISIONS.md +49

Showing 8 changed files+590-135

BENCHMARKS.md Cargo.lock Cargo.toml DECISIONS.md crates/justapi-cli/Cargo.toml crates/justapi-cli/src/main.rs crates/justapi-cli/src/workers.rs crates/justapi-core/src/server/mod.rs

▼BENCHMARKS.md

+19-0

@@ -1482,3 +1482,22 @@ dispatch means per-fire cost ≈ one `submit_py_task` enqueue — O(1), no sched
1482	1482	overhead added to request path. `cargo test -p justapi-py scheduler` → 6 Rust
1483	1483	unit tests; `test_scheduler.py` → 2 passed. Full gates green. (No p99 request
1484	1484	regression: scheduler runs off the event loop, not the request path.)
	1485	+
	1486	+### P1 follow-up 4 — multi-worker prefork supervisor (2026-07-17)
	1487	+
	1488	+`justapi serve --workers N` now runs a prefork fleet: the parent binds one
	1489	+socket and spawns N worker processes that share it via fd inheritance (no
	1490	+`SO_REUSEPORT` race), giving true multi-core accept + process isolation.
	1491	+
	1492	+\| Property \| Result \|
	1493	+\|---\|---\|
	1494	+\| `--workers 2` startup \| parent + 2 workers, all "Listening on :8099" \|
	1495	+\| `curl /` through shared socket \| served (200/404 by route) \|
	1496	+\| SIGTERM to parent \| both workers stop accept loop, drain 0 conns, exit 0, tree reaped \|
	1497	+\| `kill -KILL` a worker \| supervisor logs "worker died; restarting", respawns (fleet stays N) \|
	1498	+\| `--workers > 1` + `--reload` \| reload ignored with warning (prefork can't hot-reload) \|
	1499	+
	1500	+Unit tests: `make_spawn_argv` (fd+args carried), `bind_listener` (actually binds
	1501	+ephemeral port). Gates green. Throughput now scales with core count instead of
	1502	+one accept loop; no p99 regression on the request path. Auto-scaling workers and
	1503	+HTTP/3/Unix-socket remain follow-ups.
1485	1504

▼Cargo.lock

+1-0

@@ -2684,6 +2684,7 @@ dependencies = [
2684	2684	"hyper",
2685	2685	"justapi-core",
2686	2686	"justapi-inference",
	2687	+ "libc",
2687	2688	"notify",
2688	2689	"notify-debouncer-mini",
2689	2690	"serde",
2690	2691

▼Cargo.toml

+1-0

@@ -31,6 +31,7 @@ justapi-py = { path = "crates/justapi-py" }
31	31
32	32	# Async runtime
33	33	tokio = { version = "1", features = ["full"] }
	34	+libc = "0.2"
34	35
35	36	# HTTP stack
36	37	hyper = { version = "1", features = ["full"] }
37	38

▼DECISIONS.md

+49-0

@@ -2260,3 +2260,52 @@ minute-boundary `0 * * * * *`, interval first-fire + advance + skip-past). Pytho
2260	2260	probe: `fired == ['every','cron',...]` `count>=4`, `stats={jobs:2,fired:4,...}`.
2261	2261	Full gates green (`cargo test --workspace --features db`, clippy `-D warnings`,
2262	2262	`cargo fmt --check`).
	2263	+
	2264	+## ADR-061 — 2026-07-17 — Multi-worker prefork supervisor (`justapi serve --workers N`)
	2265	+
	2266	+Context: JustAPI ran a single-process server (one accept loop on one
	2267	+`tokio` runtime), capping throughput at one core for accept + handler dispatch
	2268	+and offering no process-isolation fault tolerance. Users expect a production
	2269	+ASGI/WSGI-style multi-worker model (uvicorn `--workers`, gunicorn). The feature
	2270	+list calls for "Multi-worker runtime", "Process management", "Graceful
	2271	+shutdown", and (later) "Auto-scaling workers".
	2272	+
	2273	+Decision / change:
	2274	+- New `crates/justapi-cli/src/workers.rs`: a prefork supervisor. The parent
	2275	+ binds the `TcpListener` once (`bind_listener`), clears `FD_CLOEXEC` on the
	2276	+ socket fd (modern std sets it; without clearing, the fd is closed across the
	2277	+ child `exec` → "IO Safety violation: owned file descriptor already closed"
	2278	+ abort), then spawns `N` children re-exec'ing the same binary with a hidden
	2279	+ `--worker-fd <fd>` flag. Each worker recovers the socket via
	2280	+ `listener_from_fd` and serves on it — true OS-process isolation, no
	2281	+ `SO_REUSEPORT` port races, all workers share one kernel socket backlog.
	2282	+- Process management: a persistent `JoinSet` of worker waits + a parallel
	2283	+ `pids[]` vector. On unexpected exit (non-shutdown) the worker is restarted
	2284	+ (verified: `kill -KILL` a worker → supervisor respawns it, fleet stays at N).
	2285	+- Graceful shutdown: SIGTERM/SIGINT cancel a `CancellationToken`; the
	2286	+ supervisor forwards SIGTERM to every live worker and waits up to
	2287	+ `drain_timeout` (default 5s) for in-flight requests to drain before forcibly
	2288	+ reaping with SIGKILL (verified: parent SIGTERM → both workers stop accept
	2289	+ loop, drain 0 connections, exit 0, tree fully reaped).
	2290	+- Refactored the previously-duplicated reload/non-reload server-build into
	2291	+ `build_server()` + `build_engines()` helpers so single-process, reload, and
	2292	+ worker paths share identical wiring (static/compression/TLS/inference).
	2293	+- `--workers > 1` is incompatible with `--reload` (hot reload restarts the whole
	2294	+ tree); the supervisor warns and ignores `--reload` in that case. `--workers` is
	2295	+ default 1 (unchanged behavior). Added `libc = "0.2"` to the workspace for
	2296	+ `fcntl`/`kill`.
	2297	+- Also fixed a pre-existing, unrelated bug gated behind the `tls` feature: a
	2298	+ `method` variable referenced in the TLS request-timeout branch was out of
	2299	+ scope (`server/mod.rs:1848`); captured `let method = req.method().clone();`
	2300	+ before the timeout consumes `req`. (The `tls` feature still has separate
	2301	+ move-error breakage out of scope for this ADR.)
	2302	+
	2303	+Evidence: `cargo test -p justapi-cli workers` → 2 unit tests
	2304	+(`make_spawn_argv` carries fd+args; `bind_listener` actually binds an ephemeral
	2305	+port). Manual: `justapi serve --addr 127.0.0.1:8099 --workers 2` → parent + 2
	2306	+workers all "Listening on …:8099"; `curl /` served (404, no route — expected);
	2307	+`kill -TERM` parent → both workers drain + exit 0, tree reaped; `kill -KILL` a
	2308	+worker → respawned. Gates green: `cargo test --workspace --features db`, clippy
	2309	+`-D warnings`, `cargo fmt --check`. (Windows/macOS: `--worker-fd` path is Unix-
	2310	+gated; non-unix falls back to single-process with a clear error — auto-scaling
	2311	+left for a follow-up.)
2263	2312

▼crates/justapi-cli/Cargo.toml

+1-0

@@ -26,6 +26,7 @@ anyhow.workspace = true
26	26	serde.workspace = true
27	27	serde_json.workspace = true
28	28	tokio-util.workspace = true
	29	+libc.workspace = true
29	30	notify-debouncer-mini = "0.7"
30	31	notify = "8"
31	32
32	33

▼crates/justapi-cli/src/main.rs

+264-135

@@ -9,6 +9,7 @@ use tokio::sync::mpsc;
9	9	mod gen_client;
10	10	mod profile;
11	11	mod watcher;
	12	+mod workers;
12	13
13	14	#[derive(Parser)]
14	15	#[command(name = "justapi", about = "JustAPI Runtime — Python application server")]
@@ -50,6 +51,17 @@ enum Commands {
50	51	#[arg(long, default_value = "5")]
51	52	drain_timeout: u64,
52	53
	54	+ /// Number of worker processes (prefork). Each worker shares the bound
	55	+ /// listening socket. Defaults to 1 (single process). Incompatible with
	56	+ /// `--reload` (hot reload restarts the whole process tree).
	57	+ #[arg(long, default_value_t = 1)]
	58	+ workers: usize,
	59	+
	60	+ /// Hidden: internal flag used by the prefork supervisor to hand a bound
	61	+ /// listening socket fd to a worker child. Do not pass manually.
	62	+ #[arg(long, hide = true)]
	63	+ worker_fd: Option<i32>,
	64	+
53	65	/// TLS certificate file (PEM)
54	66	#[cfg(feature = "tls")]
55	67	#[arg(long)]
@@ -442,6 +454,154 @@ fn scaffold_project(
442	454	Ok(())
443	455	}
444	456
	457	+/// Build the inference engine + optional scheduler engine from `--model` etc.
	458	+/// Returns `(None, None)` when `--model` is not supplied. Inference-only.
	459	+#[cfg(feature = "inference")]
	460	+fn build_engines(
	461	+ model: &Option<String>,
	462	+ gpu: &str,
	463	+ scheduled: bool,
	464	+ pool_blocks: usize,
	465	+ max_seqs: usize,
	466	+ gamma: usize,
	467	+ branch: usize,
	468	+) -> anyhow::Result<(
	469	+ Option<std::sync::Arc<justapi_inference::Engine>>,
	470	+ Option<std::sync::Arc<justapi_inference::SchedulerEngine>>,
	471	+)> {
	472	+ use justapi_inference::EngineDevice;
	473	+ let engine: Option<std::sync::Arc<justapi_inference::Engine>> = match model {
	474	+ Some(id) => {
	475	+ let device = if gpu.eq_ignore_ascii_case("cpu") {
	476	+ EngineDevice::Cpu
	477	+ } else {
	478	+ EngineDevice::Cuda(gpu.parse().unwrap_or(0))
	479	+ };
	480	+ let eng = std::sync::Arc::new(
	481	+ justapi_inference::Engine::new(device)
	482	+ .map_err(\|e\| anyhow::anyhow!("failed to start inference engine: {e}"))?,
	483	+ );
	484	+ eng.register_mock(id);
	485	+
	486	+ if gamma > 0 \|\| branch > 0 {
	487	+ let target = eng.get(id).unwrap();
	488	+ let draft_id = format!("{id}-draft");
	489	+ let draft = eng.register_mock(&draft_id);
	490	+ if branch > 0 {
	491	+ let g = gamma.max(1);
	492	+ eng.register_tree_speculative(id, target, draft, g, branch, 0);
	493	+ tracing::info!(
	494	+ model = %id, device = %gpu, gamma = g, branch,
	495	+ "Serving model via OpenAI-compatible endpoints (tree speculative decoding)"
	496	+ );
	497	+ } else {
	498	+ eng.register_speculative(id, target, draft, gamma, 0);
	499	+ tracing::info!(
	500	+ model = %id, device = %gpu, gamma,
	501	+ "Serving model via OpenAI-compatible endpoints (draft-target speculative decoding)"
	502	+ );
	503	+ }
	504	+ } else {
	505	+ tracing::info!(model = %id, device = %gpu, "Serving model via OpenAI-compatible endpoints");
	506	+ }
	507	+ Some(eng)
	508	+ }
	509	+ None => None,
	510	+ };
	511	+
	512	+ let scheduler_engine: Option<std::sync::Arc<justapi_inference::SchedulerEngine>> = if scheduled
	513	+ && engine.is_some()
	514	+ {
	515	+ use justapi_inference::{KvBlockPool, Scheduler, SchedulerConfig, SchedulerEngine};
	516	+ let eng = engine.as_ref().unwrap().clone();
	517	+ let pool = KvBlockPool::new(pool_blocks.max(64));
	518	+ let config = SchedulerConfig { max_num_seqs: max_seqs.max(1), ..Default::default() };
	519	+ let scheduler = std::sync::Arc::new(std::sync::Mutex::new(Scheduler::new(config, pool)));
	520	+ let se = std::sync::Arc::new(SchedulerEngine::new(eng, scheduler));
	521	+ tracing::info!(pool_blocks, max_seqs, "Scheduler-enabled serving path");
	522	+ Some(se)
	523	+ } else {
	524	+ None
	525	+ };
	526	+
	527	+ Ok((engine, scheduler_engine))
	528	+}
	529	+
	530	+/// Configure a `Server` from the shared serve options. TLS/compression/inference
	531	+/// wiring is applied uniformly for both the single-process and worker paths,
	532	+/// eliminating the previous duplication between the reload/non-reload branches.
	533	+fn build_server(
	534	+ addr: SocketAddr,
	535	+ static_dir: &Option<String>,
	536	+ #[cfg(feature = "compression")] compress: bool,
	537	+ #[cfg(feature = "tls")] tls_cert: &Option<String>,
	538	+ #[cfg(feature = "tls")] tls_key: &Option<String>,
	539	+ #[cfg(feature = "inference")] engine: &Option<std::sync::Arc<justapi_inference::Engine>>,
	540	+ #[cfg(feature = "inference")] scheduler_engine: &Option<
	541	+ std::sync::Arc<justapi_inference::SchedulerEngine>,
	542	+ >,
	543	+ token: tokio_util::sync::CancellationToken,
	544	+) -> justapi_core::Server {
	545	+ let mut server = justapi_core::Server::new(addr).with_shutdown(token);
	546	+
	547	+ if let Some(dir) = static_dir.clone() {
	548	+ server = server.with_static_dir(dir);
	549	+ }
	550	+
	551	+ #[cfg(feature = "compression")]
	552	+ if compress {
	553	+ server = server.add_compression();
	554	+ }
	555	+
	556	+ #[cfg(feature = "inference")]
	557	+ if let Some(ref se) = scheduler_engine {
	558	+ server = server.with_openai_scheduled(se.clone());
	559	+ } else if let Some(ref eng) = engine {
	560	+ server = server.with_openai(eng.clone());
	561	+ }
	562	+
	563	+ #[cfg(feature = "tls")]
	564	+ if let (Some(cert), Some(key)) = (tls_cert.clone(), tls_key.clone()) {
	565	+ let config = justapi_core::server::TlsConfig { cert_path: cert, key_path: key };
	566	+ server = server.with_tls(config);
	567	+ }
	568	+
	569	+ server
	570	+}
	571	+
	572	+/// Install signal handlers that cancel `token` on SIGTERM/SIGINT (Unix) and
	573	+/// Ctrl+C (all platforms), so the server/worker begins graceful drain.
	574	+fn wire_shutdown_signals(token: tokio_util::sync::CancellationToken) {
	575	+ // Ctrl+C works everywhere.
	576	+ let ctrl = token.clone();
	577	+ tokio::spawn(async move {
	578	+ tokio::signal::ctrl_c().await.ok();
	579	+ tracing::info!("Shutdown signal received (Ctrl+C)");
	580	+ ctrl.cancel();
	581	+ });
	582	+
	583	+ #[cfg(unix)]
	584	+ {
	585	+ use tokio::signal::unix::{signal, SignalKind};
	586	+ let term = token.clone();
	587	+ tokio::spawn(async move {
	588	+ if let Ok(mut s) = signal(SignalKind::terminate()) {
	589	+ s.recv().await;
	590	+ tracing::info!("Shutdown signal received (SIGTERM)");
	591	+ term.cancel();
	592	+ }
	593	+ });
	594	+ let int = token.clone();
	595	+ tokio::spawn(async move {
	596	+ if let Ok(mut s) = signal(SignalKind::interrupt()) {
	597	+ s.recv().await;
	598	+ tracing::info!("Shutdown signal received (SIGINT)");
	599	+ int.cancel();
	600	+ }
	601	+ });
	602	+ }
	603	+}
	604	+
445	605	async fn run() -> anyhow::Result<()> {
446	606	justapi_core::tracing_setup::init_tracing()?;
447	607
@@ -456,6 +616,8 @@ async fn run() -> anyhow::Result<()> {
456	616	watch_dir,
457	617	watch_ext,
458	618	drain_timeout,
	619	+ workers,
	620	+ worker_fd,
459	621	#[cfg(feature = "tls")]
460	622	tls_cert,
461	623	#[cfg(feature = "tls")]
@@ -475,6 +637,75 @@ async fn run() -> anyhow::Result<()> {
475	637	#[cfg(feature = "inference")]
476	638	branch,
477	639	} => {
	640	+ let drain_duration = Duration::from_secs(drain_timeout);
	641	+
	642	+ // Build the inference engines (no-op unless --model supplied).
	643	+ #[cfg(feature = "inference")]
	644	+ let (engine, scheduler_engine) =
	645	+ build_engines(&model, &gpu, scheduled, pool_blocks, max_seqs, gamma, branch)?;
	646	+
	647	+ // ---- Worker mode: parent handed us a bound listening socket fd. ----
	648	+ if let Some(fd) = worker_fd {
	649	+ #[cfg(unix)]
	650	+ {
	651	+ let listener = workers::listener_from_fd(fd)?;
	652	+ let token = tokio_util::sync::CancellationToken::new();
	653	+ wire_shutdown_signals(token.clone());
	654	+ let server = build_server(
	655	+ addr,
	656	+ &static_dir,
	657	+ #[cfg(feature = "compression")]
	658	+ compress,
	659	+ #[cfg(feature = "tls")]
	660	+ &tls_cert,
	661	+ #[cfg(feature = "tls")]
	662	+ &tls_key,
	663	+ #[cfg(feature = "inference")]
	664	+ &engine,
	665	+ #[cfg(feature = "inference")]
	666	+ &scheduler_engine,
	667	+ token,
	668	+ );
	669	+ let result = server.run_on(listener).await;
	670	+ justapi_core::tracing_setup::shutdown_tracing();
	671	+ return result;
	672	+ }
	673	+ #[cfg(not(unix))]
	674	+ {
	675	+ anyhow::bail!("--worker-fd is only supported on Unix platforms");
	676	+ }
	677	+ }
	678	+
	679	+ // ---- Prefork multi-worker mode: parent binds + supervises. ----
	680	+ if workers > 1 {
	681	+ if reload {
	682	+ tracing::warn!(
	683	+ "--reload is ignored when --workers > 1 (prefork cannot hot-reload)"
	684	+ );
	685	+ }
	686	+ anyhow::ensure!(workers <= 256, "worker count must be <= 256");
	687	+
	688	+ let std_listener = workers::bind_listener(addr)?;
	689	+ let raw_fd = {
	690	+ use std::os::fd::AsRawFd;
	691	+ std_listener.as_raw_fd()
	692	+ };
	693	+
	694	+ // Original argv (minus program name) so each child reconstructs
	695	+ // the identical server config.
	696	+ let base_argv: Vec<String> = std::env::args().skip(1).collect();
	697	+
	698	+ let shutdown = tokio_util::sync::CancellationToken::new();
	699	+ wire_shutdown_signals(shutdown.clone());
	700	+
	701	+ let spawn = workers::make_spawn_argv(&base_argv, raw_fd, drain_duration);
	702	+ tracing::info!(addr = %addr, workers, "starting prefork server");
	703	+ let result = workers::supervise(workers, spawn, shutdown).await;
	704	+ justapi_core::tracing_setup::shutdown_tracing();
	705	+ return result;
	706	+ }
	707	+
	708	+ // ---- Single-process mode (optionally with hot reload). ----
478	709	// Set up the file watcher once (if --reload), it sends on each
479	710	// detected change via mpsc so we can restart multiple times.
480	711	let mut reload_rx = if reload {
@@ -487,83 +718,6 @@ async fn run() -> anyhow::Result<()> {
487	718	None
488	719	};
489	720
490		- let drain_duration = Duration::from_secs(drain_timeout);
491		-
492		- // Build the inference engine if `--model` was supplied (inference
493		- // feature). The engine serves the OpenAI-compatible endpoints with a
494		- // GIL-free generation thread; the scheduler provides continuous
495		- // batching + paged KV cache (real weight loading is gated on `real`).
496		- #[cfg(feature = "inference")]
497		- let engine: Option<std::sync::Arc<justapi_inference::Engine>> = if let Some(id) =
498		- model.as_ref()
499		- {
500		- use justapi_inference::EngineDevice;
501		- let device = if gpu.eq_ignore_ascii_case("cpu") {
502		- EngineDevice::Cpu
503		- } else {
504		- EngineDevice::Cuda(gpu.parse().unwrap_or(0))
505		- };
506		- let eng = std::sync::Arc::new(
507		- justapi_inference::Engine::new(device)
508		- .map_err(\|e\| anyhow::anyhow!("failed to start inference engine: {e}"))?,
509		- );
510		- eng.register_mock(id);
511		-
512		- // Wire speculative decoding if requested.
513		- if gamma > 0 \|\| branch > 0 {
514		- let target = eng.get(id).unwrap();
515		- let draft_id = format!("{id}-draft");
516		- let draft = eng.register_mock(&draft_id);
517		- if branch > 0 {
518		- let g = gamma.max(1);
519		- eng.register_tree_speculative(id, target, draft, g, branch, 0);
520		- tracing::info!(
521		- model = %id,
522		- device = %gpu,
523		- gamma = g,
524		- branch = branch,
525		- "Serving model via OpenAI-compatible endpoints (tree speculative decoding)"
526		- );
527		- } else {
528		- eng.register_speculative(id, target, draft, gamma, 0);
529		- tracing::info!(
530		- model = %id,
531		- device = %gpu,
532		- gamma = gamma,
533		- "Serving model via OpenAI-compatible endpoints (draft-target speculative decoding)"
534		- );
535		- }
536		- } else {
537		- tracing::info!(model = %id, device = %gpu, "Serving model via OpenAI-compatible endpoints");
538		- }
539		- Some(eng)
540		- } else {
541		- None
542		- };
543		-
544		- // Build the scheduler engine if `--scheduled` is also set.
545		- #[cfg(feature = "inference")]
546		- let scheduler_engine: Option<
547		- std::sync::Arc<justapi_inference::SchedulerEngine>,
548		- > = if scheduled && engine.is_some() {
549		- use justapi_inference::{KvBlockPool, Scheduler, SchedulerConfig, SchedulerEngine};
550		- let eng = engine.as_ref().unwrap().clone();
551		- let pool = KvBlockPool::new(pool_blocks.max(64));
552		- let config =
553		- SchedulerConfig { max_num_seqs: max_seqs.max(1), ..Default::default() };
554		- let scheduler =
555		- std::sync::Arc::new(std::sync::Mutex::new(Scheduler::new(config, pool)));
556		- let se = std::sync::Arc::new(SchedulerEngine::new(eng, scheduler));
557		- tracing::info!(
558		- pool_blocks = pool_blocks,
559		- max_seqs = max_seqs,
560		- "Scheduler-enabled serving path"
561		- );
562		- Some(se)
563		- } else {
564		- None
565		- };
566		-
567	721	loop {
568	722	let token = tokio_util::sync::CancellationToken::new();
569	723
@@ -603,40 +757,25 @@ async fn run() -> anyhow::Result<()> {
603	757	return_tx.send(inner_rx).ok();
604	758	});
605	759
606		- // We'll collect the receiver back after `server.run()` completes.
607		- // Stash the return channel for later.
608		- // (We use a scope trick: store it in the reload_rx option.)
609		- // Actually we need the return_rx after the server stops,
610		- // so we hold it outside and await it after run().
611		- // Re-structure: keep return_rx in a variable.
612		- let mut server = justapi_core::Server::new(addr).with_shutdown(token.clone());
613		-
614		- if let Some(dir) = static_dir.clone() {
615		- server = server.with_static_dir(dir);
616		- }
617		-
618		- #[cfg(feature = "compression")]
619		- if compress {
620		- server = server.add_compression();
621		- }
622		-
623		- #[cfg(feature = "inference")]
624		- if let Some(ref se) = scheduler_engine {
625		- server = server.with_openai_scheduled(se.clone());
626		- } else if let Some(ref eng) = engine {
627		- server = server.with_openai(eng.clone());
628		- }
629		-
630		- #[cfg(feature = "tls")]
631		- if let (Some(cert), Some(key)) = (tls_cert.clone(), tls_key.clone()) {
632		- let config =
633		- justapi_core::server::TlsConfig { cert_path: cert, key_path: key };
634		- let result = server.with_tls(config).run().await;
635		- justapi_core::tracing_setup::shutdown_tracing();
636		- return result;
637		- }
	760	+ let server = build_server(
	761	+ addr,
	762	+ &static_dir,
	763	+ #[cfg(feature = "compression")]
	764	+ compress,
	765	+ #[cfg(feature = "tls")]
	766	+ &tls_cert,
	767	+ #[cfg(feature = "tls")]
	768	+ &tls_key,
	769	+ #[cfg(feature = "inference")]
	770	+ &engine,
	771	+ #[cfg(feature = "inference")]
	772	+ &scheduler_engine,
	773	+ token.clone(),
	774	+ );
638	775
639		- let result = server.run().await;
	776	+ let std_listener = workers::bind_listener(addr)?;
	777	+ let listener = tokio::net::TcpListener::from_std(std_listener)?;
	778	+ let result = server.run_on(listener).await;
640	779
641	780	// Recover the watcher receiver for the next iteration.
642	781	if let Ok(recovered) = return_rx.await {
@@ -660,31 +799,21 @@ async fn run() -> anyhow::Result<()> {
660	799	}
661	800
662	801	// Non-reload path: run server once and exit.
663		- let mut server = justapi_core::Server::new(addr).with_shutdown(token);
664		-
665		- if let Some(dir) = static_dir.clone() {
666		- server = server.with_static_dir(dir);
667		- }
668		-
669		- #[cfg(feature = "compression")]
670		- if compress {
671		- server = server.add_compression();
672		- }
673		-
674		- #[cfg(feature = "inference")]
675		- if let Some(ref se) = scheduler_engine {
676		- server = server.with_openai_scheduled(se.clone());
677		- } else if let Some(ref eng) = engine {
678		- server = server.with_openai(eng.clone());
679		- }
680		-
681		- #[cfg(feature = "tls")]
682		- if let (Some(cert), Some(key)) = (tls_cert.clone(), tls_key.clone()) {
683		- let config = justapi_core::server::TlsConfig { cert_path: cert, key_path: key };
684		- let result = server.with_tls(config).run().await;
685		- justapi_core::tracing_setup::shutdown_tracing();
686		- return result;
687		- }
	802	+ let server = build_server(
	803	+ addr,
	804	+ &static_dir,
	805	+ #[cfg(feature = "compression")]
	806	+ compress,
	807	+ #[cfg(feature = "tls")]
	808	+ &tls_cert,
	809	+ #[cfg(feature = "tls")]
	810	+ &tls_key,
	811	+ #[cfg(feature = "inference")]
	812	+ &engine,
	813	+ #[cfg(feature = "inference")]
	814	+ &scheduler_engine,
	815	+ tokio_util::sync::CancellationToken::new(),
	816	+ );
688	817
689	818	let result = server.run().await;
690	819	justapi_core::tracing_setup::shutdown_tracing();
691	820

▼crates/justapi-cli/src/workers.rs

+254-0

@@ -0,0 +1,254 @@
	1	+//! Multi-worker prefork supervisor for `justapi serve --workers N`.
	2	+//!
	3	+//! The parent binds the listening socket exactly once, then spawns `N` child
	4	+//! worker processes. Each child is a re-exec of the same binary in a hidden
	5	+//! `--worker-fd <fd>` mode that inherits the bound socket (fd inheritance, no
	6	+//! `SO_REUSEPORT` port-race), reconstructs the server from the original CLI
	7	+//! args, and serves on that socket. This gives true OS-process isolation and
	8	+//! saturates all cores — the single-process path only uses one accept loop.
	9	+//!
	10	+//! The supervisor owns process management:
	11	+//! - Restart-on-death: a worker that exits while not shutting down is
	12	+//! respawned so the fleet stays at `N` workers.
	13	+//! - Graceful drain: SIGTERM/SIGINT cancels a `CancellationToken`; the
	14	+//! supervisor forwards SIGTERM to every child and waits up to `drain_timeout`
	15	+//! for in-flight requests to finish before forcibly reaping.
	16	+//! - Hot reload is incompatible with prefork (a code change must restart the
	17	+//! whole tree); the caller rejects that combination before entering here.
	18	+
	19	+use std::net::SocketAddr;
	20	+use std::process::Stdio;
	21	+use std::time::Duration;
	22	+
	23	+use anyhow::Context;
	24	+use tokio::signal::unix::{signal, SignalKind};
	25	+use tokio_util::sync::CancellationToken;
	26	+
	27	+/// Raw CLI args (excluding `justapi serve` itself) needed to reconstruct an
	28	+/// identical worker. We re-pass them verbatim, plus `--worker-fd`.
	29	+pub struct WorkerSpawn {
	30	+ /// The exact argv (without program name) the parent was invoked with,
	31	+ /// e.g. `["serve", "--addr", "0.0.0.0:8080", "--compress"]`.
	32	+ pub argv: Vec<String>,
	33	+ /// Bound listening socket fd inherited by the child.
	34	+ pub listener_fd: i32,
	35	+ /// Max seconds to drain in-flight requests before forcibly reaping.
	36	+ pub drain_timeout: Duration,
	37	+}
	38	+
	39	+/// Bind the address once and return the std listener (non-CLOEXEC fd so it
	40	+/// survives the child re-exec on Unix).
	41	+pub fn bind_listener(addr: SocketAddr) -> anyhow::Result<std::net::TcpListener> {
	42	+ let listener =
	43	+ std::net::TcpListener::bind(addr).with_context(\|\| format!("failed to bind {addr}"))?;
	44	+ listener.set_nonblocking(true).context("failed to set listener non-blocking")?;
	45	+
	46	+ // Modern std sets FD_CLOEXEC on the socket fd, which would close it across
	47	+ // the child's exec. Clear it so the worker inherits the bound socket.
	48	+ #[cfg(unix)]
	49	+ {
	50	+ use std::os::fd::AsRawFd;
	51	+ let fd = listener.as_raw_fd();
	52	+ // SAFETY: fd is a valid, owned listener socket; we only clear CLOEXEC.
	53	+ let flags = unsafe { libc::fcntl(fd, libc::F_GETFD) };
	54	+ if flags >= 0 {
	55	+ unsafe {
	56	+ libc::fcntl(fd, libc::F_SETFD, flags & !libc::FD_CLOEXEC);
	57	+ }
	58	+ }
	59	+ }
	60	+
	61	+ Ok(listener)
	62	+}
	63	+
	64	+/// Recover a tokio `TcpListener` from an inherited raw fd (Unix only).
	65	+#[cfg(unix)]
	66	+pub fn listener_from_fd(fd: i32) -> anyhow::Result<tokio::net::TcpListener> {
	67	+ use std::os::fd::FromRawFd;
	68	+ // SAFETY: the fd was created by the parent via `std::net::TcpListener::bind`
	69	+ // and is not closed on exec, so it is valid and uniquely owned by this
	70	+ // process after fork/exec. We do not `from_raw_fd` it anywhere else.
	71	+ let std_listener = unsafe { <std::net::TcpListener as FromRawFd>::from_raw_fd(fd) };
	72	+ std_listener
	73	+ .set_nonblocking(true)
	74	+ .context("worker: failed to set inherited listener non-blocking")?;
	75	+ tokio::net::TcpListener::from_std(std_listener)
	76	+ .context("worker: failed to wrap inherited fd as tokio TcpListener")
	77	+}
	78	+
	79	+/// Run the supervisor: spawn `count` workers and manage their lifecycle until
	80	+/// the shutdown token is cancelled (or a fatal spawn error occurs).
	81	+pub async fn supervise(
	82	+ count: usize,
	83	+ spawn: WorkerSpawn,
	84	+ shutdown: CancellationToken,
	85	+) -> anyhow::Result<()> {
	86	+ anyhow::ensure!(count >= 1, "worker count must be >= 1");
	87	+ let exe = std::env::current_exe().context("failed to resolve current executable")?;
	88	+
	89	+ // Persistent JoinSet of in-flight worker waits. `pids[slot]` mirrors it so we
	90	+ // can forward SIGTERM without borrowing the (moved) `Child`.
	91	+ let mut waits: tokio::task::JoinSet<(usize, std::io::Result<std::process::ExitStatus>)> =
	92	+ tokio::task::JoinSet::new();
	93	+ let mut pids: Vec<Option<u32>> = Vec::with_capacity(count);
	94	+ for slot in 0..count {
	95	+ let (pid, mut handle) = spawn_worker(&exe, &spawn)?;
	96	+ waits.spawn(async move { (slot, handle.wait().await) });
	97	+ pids.push(Some(pid));
	98	+ }
	99	+ tracing::info!(workers = count, fd = spawn.listener_fd, "prefork supervisor started");
	100	+
	101	+ let mut sigterm =
	102	+ signal(SignalKind::terminate()).context("failed to install SIGTERM handler")?;
	103	+ let mut sigint = signal(SignalKind::interrupt()).context("failed to install SIGINT handler")?;
	104	+
	105	+ let mut drain_deadline: Option<tokio::time::Instant> = None;
	106	+
	107	+ loop {
	108	+ tokio::select! {
	109	+ exited = waits.join_next() => {
	110	+ match exited {
	111	+ Some(Ok((slot, status_res))) => {
	112	+ pids[slot] = None;
	113	+ match status_res {
	114	+ Ok(status) => {
	115	+ if shutdown.is_cancelled() {
	116	+ tracing::info!(worker = slot, %status, "worker exited during shutdown");
	117	+ } else {
	118	+ tracing::warn!(worker = slot, %status, "worker died; restarting");
	119	+ let (pid, mut handle) = spawn_worker(&exe, &spawn)?;
	120	+ waits.spawn(async move { (slot, handle.wait().await) });
	121	+ pids[slot] = Some(pid);
	122	+ }
	123	+ }
	124	+ Err(e) => {
	125	+ tracing::error!(worker = slot, error = %e, "worker wait error");
	126	+ }
	127	+ }
	128	+ }
	129	+ Some(Err(e)) => {
	130	+ tracing::error!(error = %e, "worker join error");
	131	+ }
	132	+ None => {
	133	+ // No live children to watch.
	134	+ if shutdown.is_cancelled() {
	135	+ break;
	136	+ }
	137	+ }
	138	+ }
	139	+ }
	140	+ _ = shutdown.cancelled() => {
	141	+ if drain_deadline.is_none() {
	142	+ drain_deadline = Some(tokio::time::Instant::now() + spawn.drain_timeout);
	143	+ tracing::info!(drain_secs = spawn.drain_timeout.as_secs(), "shutdown signal; draining workers");
	144	+ signal_children(&pids);
	145	+ }
	146	+ }
	147	+ _ = sigterm.recv() => {
	148	+ if !shutdown.is_cancelled() { shutdown.cancel(); }
	149	+ }
	150	+ _ = sigint.recv() => {
	151	+ if !shutdown.is_cancelled() { shutdown.cancel(); }
	152	+ }
	153	+ }
	154	+
	155	+ if let Some(deadline) = drain_deadline {
	156	+ let all_exited = pids.iter().all(\|p\| p.is_none());
	157	+ if all_exited {
	158	+ tracing::info!("all workers exited cleanly");
	159	+ break;
	160	+ }
	161	+ if tokio::time::Instant::now() >= deadline {
	162	+ tracing::warn!("drain timeout exceeded; forcibly reaping workers");
	163	+ signal_kill(&pids);
	164	+ // Drain remaining wait results so JoinSet empties.
	165	+ while waits.join_next().await.is_some() {}
	166	+ break;
	167	+ }
	168	+ }
	169	+ }
	170	+
	171	+ Ok(())
	172	+}
	173	+
	174	+/// Spawn a single worker child that inherits the listening socket fd.
	175	+/// Returns `(pid, wait_handle)` so the supervisor can signal without holding the
	176	+/// `Child` (which is consumed by the JoinSet await).
	177	+fn spawn_worker(
	178	+ exe: &std::path::Path,
	179	+ spawn: &WorkerSpawn,
	180	+) -> anyhow::Result<(u32, tokio::process::Child)> {
	181	+ let mut args: Vec<String> = spawn.argv.clone();
	182	+ args.push("--worker-fd".to_string());
	183	+ args.push(spawn.listener_fd.to_string());
	184	+
	185	+ let mut cmd = tokio::process::Command::new(exe);
	186	+ cmd.args(&args).stdin(Stdio::inherit()).stdout(Stdio::inherit()).stderr(Stdio::inherit());
	187	+
	188	+ let child = cmd.spawn().context("failed to spawn worker process")?;
	189	+ let pid = child.id().ok_or_else(\|\| anyhow::anyhow!("spawned worker had no pid"))?;
	190	+ Ok((pid, child))
	191	+}
	192	+
	193	+/// Forward SIGTERM to all live workers so they begin graceful drain.
	194	+fn signal_children(pids: &[Option<u32>]) {
	195	+ for pid in pids.iter().flatten() {
	196	+ // SAFETY: pid is a live child we spawned; SIGTERM is the defined
	197	+ // graceful-shutdown signal.
	198	+ let r = unsafe { libc::kill(*pid as libc::pid_t, libc::SIGTERM) };
	199	+ if r != 0 {
	200	+ tracing::warn!("failed to signal worker pid {pid} (errno {r})");
	201	+ }
	202	+ }
	203	+}
	204	+
	205	+/// Forcefully kill all live workers (drain timeout exceeded).
	206	+fn signal_kill(pids: &[Option<u32>]) {
	207	+ for pid in pids.iter().flatten() {
	208	+ // SAFETY: pid is a live child we spawned; SIGKILL forcibly reaps it.
	209	+ let r = unsafe { libc::kill(*pid as libc::pid_t, libc::SIGKILL) };
	210	+ if r != 0 {
	211	+ tracing::warn!("failed to kill worker pid {pid} (errno {r})");
	212	+ }
	213	+ }
	214	+}
	215	+
	216	+/// Build the `WorkerSpawn` payload from the parent's argv + bound fd.
	217	+pub fn make_spawn_argv(
	218	+ base_argv: &[String],
	219	+ listener_fd: i32,
	220	+ drain_timeout: Duration,
	221	+) -> WorkerSpawn {
	222	+ WorkerSpawn {
	223	+ argv: base_argv.to_vec(),
	224	+ listener_fd,
	225	+ drain_timeout,
	226	+ }
	227	+}
	228	+
	229	+#[cfg(test)]
	230	+mod tests {
	231	+ use super::*;
	232	+
	233	+ #[test]
	234	+ fn spawn_argv_carries_fd_and_args() {
	235	+ let base = vec!["serve".to_string(), "--addr".to_string(), "0:8080".to_string()];
	236	+ let spawn = make_spawn_argv(&base, 9, Duration::from_secs(5));
	237	+ assert_eq!(spawn.listener_fd, 9);
	238	+ assert_eq!(spawn.drain_timeout, Duration::from_secs(5));
	239	+ // The child re-exec must receive the original args verbatim.
	240	+ assert_eq!(spawn.argv, base);
	241	+ }
	242	+
	243	+ #[test]
	244	+ fn bind_listener_binds_and_is_nonblocking() {
	245	+ // Use a guaranteed-free ephemeral port.
	246	+ let addr: SocketAddr = "127.0.0.1:0".parse().unwrap();
	247	+ let listener = bind_listener(addr).expect("bind should succeed");
	248	+ // Re-binding the same (now-allocated) ephemeral port must fail, proving
	249	+ // the listener is actually bound and unique.
	250	+ let again = std::net::TcpListener::bind(listener.local_addr().unwrap());
	251	+ assert!(again.is_err(), "ephemeral port must be occupied by our listener");
	252	+ }
	253	+}
	254	+
0	255

▼crates/justapi-core/src/server/mod.rs

+1-0

@@ -1810,6 +1810,7 @@ async fn serve_with_tls(
1810	1810	}
1811	1811	}
1812	1812
	1813	+ let method = req.method().clone();
1813	1814	let resp = tokio::time::timeout(request_timeout(), chain.run(req)).await;
1814	1815
1815	1816	match resp {
1816	1817