tako_rs_server_pt/
lib.rs

1#![cfg_attr(docsrs, feature(doc_cfg))]
2
3//! Thread-per-core HTTP/1 server: every worker thread runs its own `SO_REUSEPORT` listener on a
4//! Tokio current-thread runtime, or compio with `compio`, and all workers share one router.
5//!
6//! ```no_run
7//! use tako_rs_server_pt::{PerThreadConfig, serve_per_thread};
8//! use tako_rs_core::router::Router;
9//!
10//! fn main() -> std::io::Result<()> {
11//!     let mut router = Router::new();
12//!     router.get("/", || async { "hello" });
13//!     let config = PerThreadConfig {
14//!         workers: 4,
15//!         pin_to_core: false,
16//!         ..Default::default()
17//!     };
18//!     serve_per_thread("0.0.0.0:8080", router, config)
19//! }
20//! ```
21
22#[cfg(not(feature = "compio"))]
23mod balance;
24mod config;
25mod listener;
26mod shutdown;
27#[cfg(not(feature = "compio"))]
28mod worker;
29#[cfg(feature = "compio")]
30mod worker_compio;
31
32use std::io;
33use std::net::SocketAddr;
34use std::str::FromStr;
35
36use tako_rs_core::router::Router;
37
38pub use crate::config::PerThreadConfig;
39pub use crate::shutdown::PerThreadShutdown;
40#[cfg(not(feature = "compio"))]
41use crate::worker::worker_main;
42#[cfg(feature = "compio")]
43use crate::worker_compio::worker_main_compio;
44
45/// Runs workers until Ctrl+C or SIGTERM, then drains requests; fails if no worker could bind.
46pub fn serve_per_thread(addr: &str, router: Router, cfg: PerThreadConfig) -> io::Result<()> {
47  let workers = cfg.workers;
48  let (handles, shutdown) = spawn_per_thread(addr, router, cfg)?;
49  wait_for_shutdown(handles, shutdown, workers)
50}
51
52/// Starts workers and returns their thread handles and the trigger that stops them.
53pub fn spawn_per_thread(
54  addr: &str,
55  router: Router,
56  cfg: PerThreadConfig,
57) -> io::Result<(Vec<std::thread::JoinHandle<()>>, PerThreadShutdown)> {
58  #[cfg(not(feature = "compio"))]
59  {
60    spawn_workers(addr, router, cfg, worker_main)
61  }
62  #[cfg(feature = "compio")]
63  {
64    spawn_workers(addr, router, cfg, worker_main_compio)
65  }
66}
67
68#[cfg(not(feature = "compio"))]
69type WorkerFn = fn(
70  usize,
71  SocketAddr,
72  std::sync::Arc<Router>,
73  PerThreadConfig,
74  PerThreadShutdown,
75  Option<balance::WorkerBalance>,
76);
77#[cfg(feature = "compio")]
78type WorkerFn = fn(usize, SocketAddr, std::sync::Arc<Router>, PerThreadConfig, PerThreadShutdown);
79
80fn spawn_workers(
81  addr: &str,
82  router: Router,
83  cfg: PerThreadConfig,
84  worker: WorkerFn,
85) -> io::Result<(Vec<std::thread::JoinHandle<()>>, PerThreadShutdown)> {
86  if cfg.workers == 0 || cfg.max_connections == Some(0) {
87    return Err(io::Error::new(
88      io::ErrorKind::InvalidInput,
89      "worker and connection counts must be positive",
90    ));
91  }
92  let addr =
93    SocketAddr::from_str(addr).map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, e))?;
94  let router = std::sync::Arc::new(router);
95  if !cfg.balance_connections || cfg!(feature = "compio") {
96    listener::warn_reuseport_platform_once();
97  }
98  let shutdown = PerThreadShutdown::new();
99  #[cfg(not(feature = "compio"))]
100  let mut balances = cfg
101    .balance_connections
102    .then(|| balance::Balancer::for_workers(cfg.workers).into_iter());
103  let mut handles = Vec::with_capacity(cfg.workers);
104  for id in 0..cfg.workers {
105    let cfg = cfg.clone();
106    let router = router.clone();
107    let worker_shutdown = shutdown.clone();
108    #[cfg(not(feature = "compio"))]
109    let balance = balances.as_mut().and_then(Iterator::next);
110    let run = move || {
111      #[cfg(not(feature = "compio"))]
112      worker(id, addr, router, cfg, worker_shutdown, balance);
113      #[cfg(feature = "compio")]
114      worker(id, addr, router, cfg, worker_shutdown);
115    };
116    match std::thread::Builder::new()
117      .name(format!("tako-pt-{id}"))
118      .spawn(run)
119    {
120      Ok(handle) => handles.push(handle),
121      Err(error) => {
122        shutdown.trigger();
123        for handle in handles {
124          let _ = handle.join();
125        }
126        return Err(error);
127      }
128    }
129  }
130  Ok((handles, shutdown))
131}
132
133fn wait_for_shutdown(
134  handles: Vec<std::thread::JoinHandle<()>>,
135  shutdown: PerThreadShutdown,
136  workers: usize,
137) -> io::Result<()> {
138  let result = tokio::runtime::Builder::new_current_thread()
139    .enable_all()
140    .build()
141    .and_then(|rt| {
142      rt.block_on(async {
143        shutdown.wait_for_bind_outcome(workers).await?;
144        tako_rs_core::server_support::shutdown_signal().await
145      })
146    });
147  shutdown.trigger();
148  for handle in handles {
149    let _ = handle.join();
150  }
151  result
152}
tako_rs_server_pt/lib.rs

tako_rs_server_pt/
lib.rs
