kino 0.4.0 → 0.6.0

data/ext/kino/src/env_strings.rs

@@ -21,6 +21,7 @@ use parking_lot::{Mutex, RwLock};
 /// These maps are probed several times per request; ahash beats the
 /// DoS-resistant default since every key here is our own static data.
 type HashMap<K, V> = std::collections::HashMap<K, V, ahash::RandomState>;
+type HashSet<K> = std::collections::HashSet<K, ahash::RandomState>;
 
 pub struct EnvStrings {
     // keys
@@ -44,17 +45,25 @@ pub struct EnvStrings {
     pub https: Opaque<RString>,
     pub http10: Opaque<RString>,
     pub http11: Opaque<RString>,
+    pub http2: Opaque<RString>,
     pub methods: HashMap<&'static str, Opaque<RString>>,
     /// lowercase header name -> frozen "HTTP_<UPPER>" key
     pub header_names: HashMap<&'static str, Opaque<RString>>,
-    /// Host-header bytes -> frozen (SERVER_NAME, SERVER_PORT) values, and
-    /// peer IP -> frozen REMOTE_ADDR value. Real traffic has low
-    /// cardinality on both, so these kill 3 string allocations per request.
+    /// Host-header or :authority bytes -> frozen host values, and peer
+    /// IP -> frozen REMOTE_ADDR value. Real traffic has low cardinality
+    /// on both, so these kill 3 string allocations per request.
     /// LRU-bounded: entries are BoxValue-rooted (registered with the GC on
     /// insert, UNregistered on eviction-drop), so a rotating-host attack
     /// recycles cache slots instead of leaking immortal strings.
-    pub hosts: Mutex<lru::LruCache<Vec<u8>, (CachedStr, CachedStr), ahash::RandomState>>,
+    pub hosts: Mutex<lru::LruCache<Vec<u8>, HostEntry, ahash::RandomState>>,
     pub addrs: Mutex<lru::LruCache<IpAddr, CachedStr, ahash::RandomState>>,
+    /// Interned values of low-cardinality headers (see
+    /// [`INTERNABLE_VALUES`]): value bytes -> frozen RString, shared
+    /// across headers that happen to carry the same bytes. Same rooting
+    /// and locking contract as `hosts`/`addrs`.
+    pub values: Mutex<lru::LruCache<Vec<u8>, CachedStr, ahash::RandomState>>,
+    /// The names whose values go through the `values` cache.
+    pub internable: HashSet<&'static str>,
     /// Ractor-shareable defaults provided by the Ruby layer at boot:
     /// the frozen rack.errors writer and the frozen null rack.input.
     pub errors_stream: RwLock<Option<Opaque<Value>>>,
@@ -63,6 +72,53 @@ pub struct EnvStrings {
 
 const HOST_CACHE_CAP: usize = 256;
 const ADDR_CACHE_CAP: usize = 1024;
+const VALUE_CACHE_CAP: usize = 512;
+
+/// Values longer than this are never interned: past it the memcpy into
+/// a fresh Ruby string is cheap relative to the bytes themselves, and
+/// unbounded keys would let one client fill the cache with garbage.
+const VALUE_INTERN_MAX_LEN: usize = 512;
+
+/// Headers whose values are effectively enums or per-install constants
+/// (a browser resends the same UA, accept-*, and sec-ch-* on every
+/// request), so caching kills an allocation + copy per header per
+/// request: the env-side analogue of what HPACK does on the wire.
+/// Deliberately absent: `cookie` and `authorization` (per-user
+/// cardinality would churn the cache, and secrets should not outlive
+/// their request in an evict-to-free cache), `referer`/`x-request-id`
+/// and friends (unbounded cardinality).
+const INTERNABLE_VALUES: &[&str] = &[
+    "user-agent",
+    "accept",
+    "accept-encoding",
+    "accept-language",
+    "cache-control",
+    "dnt",
+    "origin",
+    "pragma",
+    "priority",
+    "sec-ch-ua",
+    "sec-ch-ua-mobile",
+    "sec-ch-ua-platform",
+    "sec-fetch-dest",
+    "sec-fetch-mode",
+    "sec-fetch-site",
+    "sec-fetch-user",
+    "upgrade-insecure-requests",
+    "x-requested-with",
+];
+
+/// One hosts-cache entry: the frozen SERVER_NAME/SERVER_PORT pair, plus
+/// the frozen full authority ("host[:port]" as sent) used as the
+/// HTTP_HOST value for requests that carry the name in the URI (the h2
+/// :authority pseudo-header) rather than a Host header. Lazily filled:
+/// Host-header entries and the NUL-prefixed socket-fallback entries
+/// never allocate it.
+pub struct HostEntry {
+    name: CachedStr,
+    port: CachedStr,
+    host: Option<CachedStr>,
+}
 
 /// A frozen RString rooted via BoxValue (GC-registered address; unregisters
 /// on Drop, so LRU eviction actually frees the string).
@@ -84,6 +140,16 @@ impl CachedStr {
         CachedStr(magnus::value::BoxValue::new(string))
     }
 
+    /// Header values are bytes on the wire (not guaranteed UTF-8), so
+    /// they cache as the same binary strings `str_from_slice` builds on
+    /// the uncached path; interning must not change the encoding an
+    /// app observes.
+    fn new_from_slice(ruby: &Ruby, bytes: &[u8]) -> Self {
+        let string = ruby.str_from_slice(bytes);
+        string.freeze();
+        CachedStr(magnus::value::BoxValue::new(string))
+    }
+
     fn get(&self) -> RString {
         *self.0
     }
@@ -142,6 +208,8 @@ const COMMON_HEADERS: &[&str] = &[
     "sec-ch-ua-mobile",
     "sec-ch-ua-platform",
     "keep-alive",
+    "priority",
+    "alt-used",
 ];
 
 pub fn cgi_name(lower: &str) -> String {
@@ -194,6 +262,7 @@ pub fn init(ruby: &Ruby) {
         https: frozen(ruby, "https"),
         http10: frozen(ruby, "HTTP/1.0"),
         http11: frozen(ruby, "HTTP/1.1"),
+        http2: frozen(ruby, "HTTP/2"),
         methods,
         header_names,
         hosts: Mutex::new(lru::LruCache::with_hasher(
@@ -204,6 +273,11 @@ pub fn init(ruby: &Ruby) {
             std::num::NonZeroUsize::new(ADDR_CACHE_CAP).unwrap(),
             ahash::RandomState::new(),
         )),
+        values: Mutex::new(lru::LruCache::with_hasher(
+            std::num::NonZeroUsize::new(VALUE_CACHE_CAP).unwrap(),
+            ahash::RandomState::new(),
+        )),
+        internable: INTERNABLE_VALUES.iter().copied().collect(),
         errors_stream: RwLock::new(None),
         null_input: RwLock::new(None),
     };
@@ -249,20 +323,98 @@ pub fn set_host_env(
     let s = get();
     let mut hosts = s.hosts.lock();
     let (name, port) = match hosts.get(host) {
-        Some((name, port)) => (name.get(), port.get()),
+        Some(entry) => (entry.name.get(), entry.port.get()),
         None => {
             let (name_s, port_n) = make();
-            let entry = (
-                CachedStr::new(ruby, &name_s),
-                CachedStr::new(ruby, &port_n.to_string()),
+            let entry = HostEntry {
+                name: CachedStr::new(ruby, &name_s),
+                port: CachedStr::new(ruby, &port_n.to_string()),
+                host: None,
+            };
+            let values = (entry.name.get(), entry.port.get());
+            hosts.put(host.to_vec(), entry); // may evict + free an old entry
+            values
+        }
+    };
+    env.aset(ruby.get_inner(s.server_name), name)?;
+    env.aset(ruby.get_inner(s.server_port), port)?;
+    Ok(())
+}
+
+/// Set SERVER_NAME/SERVER_PORT *and* HTTP_HOST on `env` from the URI
+/// authority (every h2 request via :authority; also h1 absolute-form).
+/// HTTP_HOST is set here because such requests carry no Host header for
+/// the header loop to surface. Same cache and locking contract as
+/// [`set_host_env`]; an entry first created by a Host header upgrades in
+/// place, gaining the full-authority string on first use.
+pub fn set_authority_env(
+    ruby: &Ruby,
+    env: magnus::RHash,
+    authority: &str,
+    make: impl FnOnce() -> (String, u16),
+) -> Result<(), magnus::Error> {
+    let s = get();
+    let mut hosts = s.hosts.lock();
+    let (name, port, host) = match hosts.get_mut(authority.as_bytes()) {
+        Some(entry) => {
+            if entry.host.is_none() {
+                entry.host = Some(CachedStr::new(ruby, authority));
+            }
+            (
+                entry.name.get(),
+                entry.port.get(),
+                entry.host.as_ref().expect("just filled").get(),
+            )
+        }
+        None => {
+            let (name_s, port_n) = make();
+            let entry = HostEntry {
+                name: CachedStr::new(ruby, &name_s),
+                port: CachedStr::new(ruby, &port_n.to_string()),
+                host: Some(CachedStr::new(ruby, authority)),
+            };
+            let values = (
+                entry.name.get(),
+                entry.port.get(),
+                entry.host.as_ref().expect("just built").get(),
             );
-            let values = (entry.0.get(), entry.1.get());
-            hosts.put(host.to_vec(), entry); // may evict + free an old pair
+            hosts.put(authority.as_bytes().to_vec(), entry);
             values
         }
     };
     env.aset(ruby.get_inner(s.server_name), name)?;
     env.aset(ruby.get_inner(s.server_port), port)?;
+    let host_key = *s.header_names.get("host").expect("host is a common header");
+    env.aset(ruby.get_inner(host_key), host)?;
+    Ok(())
+}
+
+/// Set one header's value on `env` under `key`: through the interned
+/// value cache when the header qualifies (low-cardinality name, bounded
+/// length), else a fresh per-request string. The cached aset happens
+/// under the cache lock; see CachedStr's safety contract.
+pub fn set_value_env(
+    ruby: &Ruby,
+    env: magnus::RHash,
+    key: RString,
+    name: &str,
+    value: &[u8],
+) -> Result<(), magnus::Error> {
+    let s = get();
+    if value.len() > VALUE_INTERN_MAX_LEN || !s.internable.contains(name) {
+        return env.aset(key, ruby.str_from_slice(value));
+    }
+    let mut values = s.values.lock();
+    let cached = match values.get(value) {
+        Some(cached) => cached.get(),
+        None => {
+            let entry = CachedStr::new_from_slice(ruby, value);
+            let string = entry.get();
+            values.put(value.to_vec(), entry); // may evict + free an old value
+            string
+        }
+    };
+    env.aset(key, cached)?;
     Ok(())
 }
 
@@ -302,7 +454,11 @@ mod tests {
     fn common_headers_are_lowercase_and_unique() {
         let mut seen = std::collections::HashSet::new();
         for header in COMMON_HEADERS {
-            assert_eq!(*header, header.to_ascii_lowercase(), "{header} must be lowercase");
+            assert_eq!(
+                *header,
+                header.to_ascii_lowercase(),
+                "{header} must be lowercase"
+            );
             assert!(seen.insert(*header), "{header} listed twice");
         }
     }
@@ -311,7 +467,11 @@ mod tests {
     fn common_methods_are_uppercase_and_unique() {
         let mut seen = std::collections::HashSet::new();
         for method in COMMON_METHODS {
-            assert_eq!(*method, method.to_ascii_uppercase(), "{method} must be uppercase");
+            assert_eq!(
+                *method,
+                method.to_ascii_uppercase(),
+                "{method} must be uppercase"
+            );
             assert!(seen.insert(*method), "{method} listed twice");
         }
     }

data/ext/kino/src/io_shards.rs

@@ -0,0 +1,437 @@
+//! Current-thread Tokio runtimes for HTTP I/O.
+//!
+//! One accept thread owns the listener and assigns accepted connections to
+//! the least-loaded shard. Each shard then owns that connection for its
+//! lifetime, avoiding the shared Tokio worker pool on hot HTTP paths.
+
+use std::net::SocketAddr;
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::Arc;
+use std::thread::JoinHandle;
+
+use crate::listen::Listener;
+use crate::log::{self, Level};
+use crate::registry::{ServerInner, STATE_DRAINING};
+use crate::server::{serve_conn, AsyncListener, Conn};
+
+/// A connection in transit from the acceptor to its shard. Tokio streams
+/// are bound to the runtime that registered them, so the handoff carries
+/// the std stream and the shard re-registers it on arrival.
+enum StdConn {
+    Tcp(std::net::TcpStream),
+    Unix(std::os::unix::net::UnixStream),
+}
+
+impl StdConn {
+    /// Register with the calling (shard) runtime.
+    fn into_tokio(self) -> std::io::Result<Conn> {
+        Ok(match self {
+            StdConn::Tcp(stream) => Conn::Tcp(tokio::net::TcpStream::from_std(stream)?),
+            StdConn::Unix(stream) => Conn::Unix(tokio::net::UnixStream::from_std(stream)?),
+        })
+    }
+}
+
+/// Detach an accepted stream from the acceptor's runtime for the handoff.
+fn into_std(conn: Conn) -> std::io::Result<StdConn> {
+    Ok(match conn {
+        Conn::Tcp(stream) => StdConn::Tcp(stream.into_std()?),
+        Conn::Unix(stream) => StdConn::Unix(stream.into_std()?),
+    })
+}
+
+/// One accepted connection en route to a shard: the detached stream, the
+/// addresses hyper reports, and the slot it holds against max_connections.
+struct Accepted {
+    conn: StdConn,
+    remote_addr: SocketAddr,
+    local_addr: SocketAddr,
+    permit: tokio::sync::OwnedSemaphorePermit,
+}
+
+/// Shard count: an explicit `io_threads` wins; the default is half the
+/// available CPUs. Framing requests is cheap next to running the app, so
+/// the I/O plane gets the smaller share and Ruby workers keep the rest.
+pub(crate) fn thread_count(io_threads: usize) -> usize {
+    if io_threads > 0 {
+        return io_threads;
+    }
+    default_thread_count(std::thread::available_parallelism().map_or(1, usize::from))
+}
+
+fn default_thread_count(cpus: usize) -> usize {
+    cpus.div_ceil(2)
+}
+
+/// Boot the shard threads, then the acceptor. Any thread that fails to
+/// come up fails the whole boot: the already started threads are drained
+/// (their senders drop) and joined before the error reaches Ruby.
+pub(crate) fn spawn(
+    listener: Listener,
+    acceptor: Option<tokio_rustls::TlsAcceptor>,
+    server: Arc<ServerInner>,
+    max_connections: usize,
+    accept_shutdown_rx: tokio::sync::watch::Receiver<bool>,
+    runtime_shutdown_rx: tokio::sync::watch::Receiver<bool>,
+    shard_count: usize,
+) -> std::io::Result<Vec<JoinHandle<()>>> {
+    let shard_count = shard_count.max(1);
+    let mut handles = Vec::with_capacity(shard_count + 1);
+    let mut shard_txs = Vec::with_capacity(shard_count);
+    let mut loads = Vec::with_capacity(shard_count);
+
+    for i in 0..shard_count {
+        let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
+        let load = Arc::new(AtomicUsize::new(0));
+        let spawned = spawn_shard(
+            i,
+            rx,
+            acceptor.clone(),
+            server.clone(),
+            load.clone(),
+            runtime_shutdown_rx.clone(),
+        );
+        match spawned {
+            Ok(handle) => {
+                shard_txs.push(tx);
+                loads.push(load);
+                handles.push(handle);
+            }
+            Err(error) => {
+                drop(tx);
+                drop(shard_txs);
+                join_all(handles);
+                return Err(error);
+            }
+        }
+    }
+
+    match spawn_acceptor(
+        listener,
+        server,
+        max_connections,
+        accept_shutdown_rx,
+        shard_txs,
+        loads,
+    ) {
+        Ok(handle) => handles.push(handle),
+        Err(error) => {
+            join_all(handles);
+            return Err(error);
+        }
+    }
+    Ok(handles)
+}
+
+fn join_all(handles: Vec<JoinHandle<()>>) {
+    for handle in handles {
+        let _ = handle.join();
+    }
+}
+
+fn current_thread_runtime() -> std::io::Result<tokio::runtime::Runtime> {
+    tokio::runtime::Builder::new_current_thread()
+        .enable_all()
+        .build()
+}
+
+/// Wait for a just spawned I/O thread to report its runtime up, so a
+/// startup failure becomes the boot error Ruby sees instead of a silently
+/// dead thread.
+fn await_ready(
+    handle: JoinHandle<()>,
+    ready_rx: std::sync::mpsc::Receiver<std::io::Result<()>>,
+    what: &str,
+) -> std::io::Result<JoinHandle<()>> {
+    match ready_rx.recv() {
+        Ok(Ok(())) => Ok(handle),
+        Ok(Err(error)) => {
+            let _ = handle.join();
+            Err(error)
+        }
+        Err(_) => Err(std::io::Error::other(format!(
+            "{what} thread exited during startup"
+        ))),
+    }
+}
+
+fn spawn_shard(
+    index: usize,
+    rx: tokio::sync::mpsc::UnboundedReceiver<Accepted>,
+    acceptor: Option<tokio_rustls::TlsAcceptor>,
+    server: Arc<ServerInner>,
+    load: Arc<AtomicUsize>,
+    mut shutdown_rx: tokio::sync::watch::Receiver<bool>,
+) -> std::io::Result<JoinHandle<()>> {
+    let (ready_tx, ready_rx) = std::sync::mpsc::sync_channel(1);
+    let handle = std::thread::Builder::new()
+        .name(format!("kino-io-{index}"))
+        .spawn(move || {
+            let runtime = match current_thread_runtime() {
+                Ok(runtime) => runtime,
+                Err(error) => {
+                    let _ = ready_tx.send(Err(error));
+                    return;
+                }
+            };
+            let _ = ready_tx.send(Ok(()));
+            runtime.block_on(shard_loop(rx, acceptor, server, load, &mut shutdown_rx));
+        })?;
+    await_ready(handle, ready_rx, "shard")
+}
+
+fn spawn_acceptor(
+    listener: Listener,
+    server: Arc<ServerInner>,
+    max_connections: usize,
+    shutdown_rx: tokio::sync::watch::Receiver<bool>,
+    shard_txs: Vec<tokio::sync::mpsc::UnboundedSender<Accepted>>,
+    loads: Vec<Arc<AtomicUsize>>,
+) -> std::io::Result<JoinHandle<()>> {
+    let (ready_tx, ready_rx) = std::sync::mpsc::sync_channel(1);
+    let handle = std::thread::Builder::new()
+        .name("kino-accept".to_string())
+        .spawn(move || {
+            let runtime = match current_thread_runtime() {
+                Ok(runtime) => runtime,
+                Err(error) => {
+                    let _ = ready_tx.send(Err(error));
+                    return;
+                }
+            };
+            runtime.block_on(async move {
+                // Registration must happen on this runtime; a failure is
+                // routed through the same ready channel as a build error.
+                let listener = match AsyncListener::from_std(listener) {
+                    Ok(listener) => listener,
+                    Err(error) => {
+                        let _ = ready_tx.send(Err(error));
+                        return;
+                    }
+                };
+                accept_loop(
+                    listener,
+                    server,
+                    max_connections,
+                    shutdown_rx,
+                    shard_txs,
+                    loads,
+                    ready_tx,
+                )
+                .await;
+            });
+        })?;
+    await_ready(handle, ready_rx, "accept")
+}
+
+/// Serve handed-over connections until the acceptor is gone, then let the
+/// remaining ones finish. The final teardown signal cuts either phase
+/// short: dropping the runtime cancels connection tasks at their next
+/// await point.
+async fn shard_loop(
+    mut rx: tokio::sync::mpsc::UnboundedReceiver<Accepted>,
+    acceptor: Option<tokio_rustls::TlsAcceptor>,
+    server: Arc<ServerInner>,
+    load: Arc<AtomicUsize>,
+    shutdown_rx: &mut tokio::sync::watch::Receiver<bool>,
+) {
+    let mut connections = tokio::task::JoinSet::new();
+    loop {
+        tokio::select! {
+            _ = shutdown_rx.changed() => return,
+            accepted = rx.recv() => {
+                let Some(accepted) = accepted else { break };
+                let acceptor = acceptor.clone();
+                let server = server.clone();
+                let guard = LoadGuard(load.clone());
+                connections.spawn(async move {
+                    let _guard = guard;
+                    serve_accepted(accepted, acceptor, server).await;
+                });
+            }
+            // Reap closed connections so the set stays small.
+            Some(_) = connections.join_next(), if !connections.is_empty() => {}
+        }
+    }
+    // The acceptor is gone: drain. Every join is one connection closing.
+    while !connections.is_empty() {
+        tokio::select! {
+            _ = shutdown_rx.changed() => return,
+            _ = connections.join_next() => {}
+        }
+    }
+}
+
+/// Keeps the shard's connection count honest whichever way the task ends:
+/// return, panic, or cancellation at teardown. A plain decrement after the
+/// await would never run on the last two.
+struct LoadGuard(Arc<AtomicUsize>);
+
+impl Drop for LoadGuard {
+    fn drop(&mut self) {
+        self.0.fetch_sub(1, Ordering::Relaxed);
+    }
+}
+
+/// The shard's half of the handoff: re-register the stream on this
+/// runtime, then run the shared connection pipeline (TLS handshake,
+/// protocol layer) exactly as the default runtime would.
+async fn serve_accepted(
+    accepted: Accepted,
+    acceptor: Option<tokio_rustls::TlsAcceptor>,
+    server: Arc<ServerInner>,
+) {
+    // Held for the connection's lifetime; dropping it frees a slot.
+    let _permit = accepted.permit;
+    let conn = match accepted.conn.into_tokio() {
+        Ok(conn) => conn,
+        Err(_) => {
+            log::emit(
+                Level::Warn,
+                "tokio",
+                "failed to register a stream on an I/O shard",
+            );
+            return;
+        }
+    };
+    serve_conn(
+        conn,
+        acceptor,
+        server,
+        accepted.remote_addr,
+        accepted.local_addr,
+    )
+    .await;
+}
+
+/// The sharded accept loop. Same backpressure as the default loop: the
+/// permit is acquired before accept, so past max_connections the excess
+/// waits in the kernel backlog instead of being accepted and dropped.
+/// A shard whose channel is gone is marked dead and routed around; with
+/// no shard left the loop stops accepting and flips the server to
+/// draining, so the control plane stops reporting ready.
+async fn accept_loop(
+    listener: AsyncListener,
+    server: Arc<ServerInner>,
+    max_connections: usize,
+    mut shutdown_rx: tokio::sync::watch::Receiver<bool>,
+    shard_txs: Vec<tokio::sync::mpsc::UnboundedSender<Accepted>>,
+    loads: Vec<Arc<AtomicUsize>>,
+    ready_tx: std::sync::mpsc::SyncSender<std::io::Result<()>>,
+) {
+    let conn_limit = Arc::new(tokio::sync::Semaphore::new(max_connections));
+    let mut live = vec![true; shard_txs.len()];
+    let _ = ready_tx.send(Ok(()));
+    'accept: loop {
+        let permit = tokio::select! {
+            _ = shutdown_rx.changed() => break,
+            permit = conn_limit.clone().acquire_owned() => match permit {
+                Ok(permit) => permit,
+                Err(_) => break,
+            },
+        };
+        let (conn, remote_addr, local_addr) = tokio::select! {
+            _ = shutdown_rx.changed() => break,
+            accepted = listener.accept() => match accepted {
+                Ok(accepted) => accepted,
+                Err(_) => continue,
+            },
+        };
+        let conn = match into_std(conn) {
+            Ok(conn) => conn,
+            Err(_) => {
+                log::emit(
+                    Level::Warn,
+                    "tokio",
+                    "failed to detach an accepted stream; connection dropped",
+                );
+                continue;
+            }
+        };
+        let mut accepted = Accepted {
+            conn,
+            remote_addr,
+            local_addr,
+            permit,
+        };
+        loop {
+            let Some(index) = least_loaded(&loads, &live) else {
+                log::emit(
+                    Level::Error,
+                    "tokio",
+                    "all I/O shards are down; not accepting connections",
+                );
+                server.state.store(STATE_DRAINING, Ordering::Relaxed);
+                break 'accept;
+            };
+            loads[index].fetch_add(1, Ordering::Relaxed);
+            match shard_txs[index].send(accepted) {
+                Ok(()) => break,
+                Err(returned) => {
+                    loads[index].fetch_sub(1, Ordering::Relaxed);
+                    live[index] = false;
+                    log::emit(Level::Warn, "tokio", "I/O shard is down; routing around it");
+                    accepted = returned.0;
+                }
+            }
+        }
+    }
+}
+
+/// The live shard with the fewest open connections.
+fn least_loaded(loads: &[Arc<AtomicUsize>], live: &[bool]) -> Option<usize> {
+    loads
+        .iter()
+        .enumerate()
+        .filter(|(index, _)| live[*index])
+        .min_by_key(|(_, load)| load.load(Ordering::Relaxed))
+        .map(|(index, _)| index)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::{default_thread_count, least_loaded, thread_count, Arc, AtomicUsize};
+
+    fn loads(counts: &[usize]) -> Vec<Arc<AtomicUsize>> {
+        counts
+            .iter()
+            .map(|&n| Arc::new(AtomicUsize::new(n)))
+            .collect()
+    }
+
+    #[test]
+    fn explicit_io_threads_win() {
+        assert_eq!(thread_count(3), 3);
+    }
+
+    #[test]
+    fn least_loaded_picks_the_emptiest_live_shard() {
+        assert_eq!(
+            least_loaded(&loads(&[3, 0, 1]), &[true, true, true]),
+            Some(1)
+        );
+    }
+
+    #[test]
+    fn least_loaded_routes_around_dead_shards() {
+        // The dead shard's count is frozen at 0; it must not win anyway.
+        assert_eq!(
+            least_loaded(&loads(&[3, 0, 1]), &[true, false, true]),
+            Some(2)
+        );
+    }
+
+    #[test]
+    fn least_loaded_reports_when_no_shard_is_left() {
+        assert_eq!(least_loaded(&loads(&[0, 0]), &[false, false]), None);
+    }
+
+    #[test]
+    fn default_is_half_the_cpus() {
+        assert_eq!(default_thread_count(1), 1);
+        assert_eq!(default_thread_count(2), 1);
+        assert_eq!(default_thread_count(3), 2);
+        assert_eq!(default_thread_count(12), 6);
+        assert_eq!(default_thread_count(128), 64);
+    }
+}