kino 0.4.0 → 0.6.0

data/ext/kino/src/lib.rs

@@ -9,6 +9,7 @@ mod control;
 mod cpus;
 mod env_strings;
 mod gvl;
+mod io_shards;
 mod listen;
 mod log;
 mod logsink;
@@ -40,8 +41,7 @@ fn init(ruby: &Ruby) -> Result<(), Error> {
     let native = module.define_module("Native")?;
     native.define_singleton_method("server_start", function!(server::server_start, 1))?;
     native.define_singleton_method("register_worker", function!(server::register_worker, 1))?;
-    native.define_singleton_method("take_one", function!(queue::take_one, 2))?;
-    native.define_singleton_method("take_batch", function!(queue::take_batch, 3))?;
+    native.define_singleton_method("worker", function!(queue::worker, 2))?;
     native.define_singleton_method("stop_accepting", function!(server::stop_accepting, 1))?;
     native.define_singleton_method("close_queue", function!(server::close_queue, 1))?;
     native.define_singleton_method("queue_stats", function!(server::queue_stats, 1))?;
@@ -83,11 +83,15 @@ fn init(ruby: &Ruby) -> Result<(), Error> {
     native.define_class("PinKeeper", ruby.class_object())?;
     native.define_singleton_method("pin_keeper", function!(server::pin_keeper, 1))?;
 
+    let worker = native.define_class("Worker", ruby.class_object())?;
+    worker.define_method("take_one", method!(queue::Worker::take_one, 0))?;
+    worker.define_method("take_batch", method!(queue::Worker::take_batch, 1))?;
+
     let request = native.define_class("Request", ruby.class_object())?;
-    request.define_method("respond_and_take", method!(queue::respond_and_take, 6))?;
+    request.define_method("respond_and_take", method!(queue::respond_and_take, 5))?;
     request.define_method(
         "respond_and_take_one",
-        method!(queue::respond_and_take_one, 5),
+        method!(queue::respond_and_take_one, 4),
     )?;
     request.define_method("read_body", method!(Request::read_body, 1))?;
     request.define_method("send_simple", method!(crate::request::respond_simple, 3))?;
@@ -97,10 +101,11 @@ fn init(ruby: &Ruby) -> Result<(), Error> {
     request.define_method("abort", method!(Request::abort, 0))?;
     request.define_method("timing", method!(Request::set_timing, 2))?;
 
-    // Force-resolve the TypedData class cache on the main ractor: magnus
-    // resolves it lazily on first wrap, and a racy first resolution from two
-    // worker ractors is the failure mode we must rule out.
+    // Force-resolve the TypedData class caches on the main ractor: magnus
+    // resolves them lazily on first wrap, and a racy first resolution from
+    // two worker ractors is the failure mode we must rule out.
     let _ = <Request as magnus::TypedData>::class(ruby);
+    let _ = <queue::Worker as magnus::TypedData>::class(ruby);
 
     // Frozen env key/value cache: built once here (main ractor, GVL held),
     // shared by every worker ractor afterwards.

data/ext/kino/src/listen.rs

@@ -86,7 +86,10 @@ mod tests {
 
     #[test]
     fn unix_path_recognises_only_the_unix_scheme() {
-        assert_eq!(unix_path("unix:///run/kino.sock").unwrap().to_str(), Some("/run/kino.sock"));
+        assert_eq!(
+            unix_path("unix:///run/kino.sock").unwrap().to_str(),
+            Some("/run/kino.sock")
+        );
         assert!(unix_path("127.0.0.1").is_none());
         assert!(unix_path("unix.example.com").is_none());
     }

data/ext/kino/src/log.rs

@@ -99,7 +99,10 @@ mod tests {
     #[test]
     fn label_is_a_syslog_tag_plus_the_source() {
         assert_eq!(label(4213, "main"), "kino[4213] main:");
-        assert_eq!(label(4213, "worker-3/thread-2"), "kino[4213] worker-3/thread-2:");
+        assert_eq!(
+            label(4213, "worker-3/thread-2"),
+            "kino[4213] worker-3/thread-2:"
+        );
     }
 
     #[test]
@@ -129,7 +132,12 @@ mod tests {
     #[test]
     fn a_report_is_labelled_on_its_first_line_only() {
         assert_eq!(
-            format_line(Level::Error, "kino[1] main:", "500 GET / · X: y\n    a.rb:1", false),
+            format_line(
+                Level::Error,
+                "kino[1] main:",
+                "500 GET / · X: y\n    a.rb:1",
+                false
+            ),
             "kino[1] main: 500 GET / · X: y\n    a.rb:1"
         );
     }

data/ext/kino/src/queue.rs

@@ -117,7 +117,8 @@ fn admit(
 ) -> Result<RHash, Error> {
     server.served.fetch_add(1, Ordering::Relaxed);
     slot.served.fetch_add(1, Ordering::Relaxed);
-    slot.last_started_ms.store(crate::mono::mono_ms(), Ordering::Relaxed);
+    slot.last_started_ms
+        .store(crate::mono::mono_ms(), Ordering::Relaxed);
     slot.in_flight.fetch_add(1, Ordering::Relaxed);
     // One clock read serves the histogram and, for the access log, the
     // request's queue wait and the start of its time in Ruby.
@@ -137,54 +138,73 @@ fn admit(
     Ok(env)
 }
 
-type Checkout = (Arc<ServerInner>, Arc<WorkerSlot>, BoxedCtx);
+/// Per-worker native handle: the registry lookup and slot resolution are
+/// paid once at worker boot instead of on every take. Created inside the
+/// worker thread/ractor that uses it (like Request handles), so ownership
+/// is correct by construction. Holds the server weakly: after teardown the
+/// upgrade fails, which is the same clean shutdown signal the per-take
+/// registry lookup used to give.
+#[magnus::wrap(class = "Kino::Native::Worker", free_immediately)]
+pub struct Worker {
+    server: std::sync::Weak<ServerInner>,
+    slot: Arc<WorkerSlot>,
+}
 
-fn checkout(ruby: &Ruby, server_id: u64, worker_id: usize) -> Result<Option<Checkout>, Error> {
+/// Resolve a (server, worker) pair into a Worker handle; nil when the
+/// server is already gone (the caller treats that as shutdown).
+pub fn worker(ruby: &Ruby, server_id: u64, worker_id: usize) -> Result<Option<Worker>, Error> {
     let Some(server) = registry::try_get(server_id) else {
-        return Ok(None); // server torn down → clean shutdown signal
+        return Ok(None);
     };
     let slot = server.slot(ruby, worker_id)?;
+    Ok(Some(Worker {
+        server: Arc::downgrade(&server),
+        slot,
+    }))
+}
 
-    // The previous batch is fully answered once the worker comes back.
-    slot.current.lock().clear();
-    slot.in_flight.store(0, Ordering::Relaxed);
-    slot.interrupted.store(false, Ordering::SeqCst);
+impl Worker {
+    fn checkout(&self) -> Result<Option<(Arc<ServerInner>, BoxedCtx)>, Error> {
+        let Some(server) = self.server.upgrade() else {
+            return Ok(None); // server torn down → clean shutdown signal
+        };
 
-    Ok(block_take(&server, &slot)?.map(|ctx| (server, slot, ctx)))
-}
+        // The previous batch is fully answered once the worker comes back.
+        self.slot.current.lock().clear();
+        self.slot.in_flight.store(0, Ordering::Relaxed);
+        self.slot.interrupted.store(false, Ordering::SeqCst);
 
-/// Take one request; returns its env Hash (request handle inside under
-/// "kino.request") or nil on shutdown. The batch-of-one hot path: no
-/// arrays allocated at all.
-pub fn take_one(ruby: &Ruby, server_id: u64, worker_id: usize) -> Result<Option<RHash>, Error> {
-    match checkout(ruby, server_id, worker_id)? {
-        Some((server, slot, ctx)) => Ok(Some(admit(ruby, &server, &slot, ctx)?)),
-        None => Ok(None),
+        Ok(block_take(&server, &self.slot)?.map(|ctx| (server, ctx)))
     }
-}
 
-/// Take up to `max` requests: block for the first, drain the rest
-/// non-blocking (they only batch when the queue is already deep).
-/// Returns nil on shutdown; otherwise an Array of env Hashes.
-pub fn take_batch(
-    ruby: &Ruby,
-    server_id: u64,
-    worker_id: usize,
-    max: usize,
-) -> Result<Option<RArray>, Error> {
-    let Some((server, slot, first)) = checkout(ruby, server_id, worker_id)? else {
-        return Ok(None);
-    };
+    /// Take one request; returns its env Hash (request handle inside under
+    /// "kino.request") or nil on shutdown. The batch-of-one hot path: no
+    /// arrays allocated at all.
+    pub fn take_one(ruby: &Ruby, rb_self: &Worker) -> Result<Option<RHash>, Error> {
+        match rb_self.checkout()? {
+            Some((server, ctx)) => Ok(Some(admit(ruby, &server, &rb_self.slot, ctx)?)),
+            None => Ok(None),
+        }
+    }
 
-    let batch = ruby.ary_new_capa(max.max(1));
-    batch.push(admit(ruby, &server, &slot, first)?)?;
-    for _ in 1..max {
-        match server.req_rx.try_recv() {
-            Ok(ctx) => batch.push(admit(ruby, &server, &slot, ctx)?)?,
-            Err(_) => break,
+    /// Take up to `max` requests: block for the first, drain the rest
+    /// non-blocking (they only batch when the queue is already deep).
+    /// Returns nil on shutdown; otherwise an Array of env Hashes.
+    pub fn take_batch(ruby: &Ruby, rb_self: &Worker, max: usize) -> Result<Option<RArray>, Error> {
+        let Some((server, first)) = rb_self.checkout()? else {
+            return Ok(None);
+        };
+
+        let batch = ruby.ary_new_capa(max.max(1));
+        batch.push(admit(ruby, &server, &rb_self.slot, first)?)?;
+        for _ in 1..max {
+            match server.req_rx.try_recv() {
+                Ok(ctx) => batch.push(admit(ruby, &server, &rb_self.slot, ctx)?)?,
+                Err(_) => break,
+            }
         }
+        Ok(Some(batch))
     }
-    Ok(Some(batch))
 }
 
 /// The fused hot path: answer `request` (complete response in one shot)
@@ -193,28 +213,25 @@ pub fn take_batch(
 pub fn respond_and_take_one(
     ruby: &Ruby,
     request: &Request,
-    server_id: u64,
-    worker_id: usize,
+    worker: magnus::typed_data::Obj<Worker>,
     status: u16,
     headers: RHash,
     body: RString,
 ) -> Result<Option<RHash>, Error> {
     crate::request::respond_simple(ruby, request, status, headers, body)?;
-    take_one(ruby, server_id, worker_id)
+    Worker::take_one(ruby, &worker)
 }
 
 /// Batch variant of the fused call.
-#[allow(clippy::too_many_arguments)]
 pub fn respond_and_take(
     ruby: &Ruby,
     request: &Request,
-    server_id: u64,
-    worker_id: usize,
+    worker: magnus::typed_data::Obj<Worker>,
     max: usize,
     status: u16,
     headers: RHash,
     body: RString,
 ) -> Result<Option<RArray>, Error> {
     crate::request::respond_simple(ruby, request, status, headers, body)?;
-    take_batch(ruby, server_id, worker_id, max)
+    Worker::take_batch(ruby, &worker, max)
 }

data/ext/kino/src/registry.rs

@@ -5,6 +5,7 @@
 
 use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
 use std::sync::{Arc, OnceLock, Weak};
+use std::time::{Duration, Instant};
 
 use parking_lot::{Mutex, RwLock};
 
@@ -32,6 +33,51 @@ pub type BoxedCtx = Box<RequestCtx>;
 /// Probed on every take; keys are our own ids, so ahash over SipHash.
 type HashMap<K, V> = std::collections::HashMap<K, V, ahash::RandomState>;
 
+/// What runs the server's I/O, taken out at shutdown. The default
+/// multi-thread runtime and sharded I/O carry different teardown state,
+/// so the shape stays explicit instead of parallel optional fields.
+#[derive(Default)]
+pub enum RuntimeHandle {
+    /// Not started yet, or already shut down.
+    #[default]
+    None,
+    MultiThread(tokio::runtime::Runtime),
+    /// The shards' final-teardown signal and every I/O thread (shards plus
+    /// the acceptor).
+    Shards {
+        shutdown_tx: tokio::sync::watch::Sender<bool>,
+        threads: Vec<std::thread::JoinHandle<()>>,
+    },
+}
+
+impl RuntimeHandle {
+    /// Stop the I/O side, giving in-flight work `timeout` to finish. A
+    /// thread that misses the deadline is abandoned rather than joined:
+    /// a wedged shard must not hang `Server#shutdown`, which promises to
+    /// return by its deadline.
+    pub fn shutdown(self, timeout: Duration) {
+        match self {
+            RuntimeHandle::None => {}
+            RuntimeHandle::MultiThread(runtime) => runtime.shutdown_timeout(timeout),
+            RuntimeHandle::Shards {
+                shutdown_tx,
+                threads,
+            } => {
+                let _ = shutdown_tx.send(true);
+                let deadline = Instant::now() + timeout;
+                for thread in threads {
+                    while !thread.is_finished() && Instant::now() < deadline {
+                        std::thread::sleep(Duration::from_millis(1));
+                    }
+                    if thread.is_finished() {
+                        let _ = thread.join();
+                    }
+                }
+            }
+        }
+    }
+}
+
 /// One per `Kino::Server`. Owns the tokio runtime, the request queue and the
 /// worker slots.
 pub struct ServerInner {
@@ -42,9 +88,9 @@ pub struct ServerInner {
     pub req_rx: flume::Receiver<BoxedCtx>,
     /// Signals the accept loop to stop. Watch channel: `true` = draining.
     pub shutdown_tx: tokio::sync::watch::Sender<bool>,
-    /// Runtime is kept so we can shut it down explicitly; in an Option so
-    /// `shutdown_runtime` can take ownership out of the Arc.
-    pub runtime: Mutex<Option<tokio::runtime::Runtime>>,
+    /// Runtime is kept so we can shut it down explicitly; `shutdown_runtime`
+    /// takes ownership out of the Arc.
+    pub runtime: Mutex<RuntimeHandle>,
     pub slots: RwLock<Vec<Arc<WorkerSlot>>>,
     pub in_flight: AtomicUsize,
     /// Requests handed to Ruby workers (admitted), and requests rejected
@@ -69,6 +115,9 @@ pub struct ServerInner {
     pub quarantine_replacements: AtomicU64,
     pub topology: Topology,
     pub https: bool,
+    /// HTTP/2 serving (ALPN over TLS, prior-knowledge h2c on plaintext);
+    /// false pins every connection to HTTP/1.
+    pub http2: bool,
     /// The socket file of a `unix://` bind, removed at shutdown.
     pub unix_path: Option<std::path::PathBuf>,
     /// Native access log sink (None unless log_requests is on).
@@ -118,8 +167,8 @@ pub const LANE_DEPTH: usize = 4;
 /// Fixed queue-wait bucket boundaries in microseconds (0.5ms .. 10s),
 /// ascending. Emitted in seconds. Not a knob (YAGNI).
 pub const QUEUE_BOUNDS_US: [u64; 14] = [
-    500, 1_000, 2_500, 5_000, 10_000, 25_000, 50_000, 100_000,
-    250_000, 500_000, 1_000_000, 2_500_000, 5_000_000, 10_000_000,
+    500, 1_000, 2_500, 5_000, 10_000, 25_000, 50_000, 100_000, 250_000, 500_000, 1_000_000,
+    2_500_000, 5_000_000, 10_000_000,
 ];
 
 /// Queue-wait histogram: per-bucket counts plus an overflow (the implicit
@@ -289,7 +338,7 @@ pub fn test_server(lanes: bool, queue_depth: usize) -> Arc<ServerInner> {
         req_tx: Mutex::new(Some(req_tx)),
         req_rx,
         shutdown_tx,
-        runtime: Mutex::new(None),
+        runtime: Mutex::new(RuntimeHandle::None),
         slots: RwLock::new(Vec::new()),
         in_flight: AtomicUsize::new(0),
         served: AtomicU64::new(0),
@@ -301,8 +350,14 @@ pub fn test_server(lanes: bool, queue_depth: usize) -> Arc<ServerInner> {
         state: std::sync::atomic::AtomicU8::new(STATE_BOOTING),
         respawns: AtomicU64::new(0),
         quarantine_replacements: AtomicU64::new(0),
-        topology: Topology { mode: "threaded".to_string(), workers: 0, threads: 0, batch: 1 },
+        topology: Topology {
+            mode: "threaded".to_string(),
+            workers: 0,
+            threads: 0,
+            batch: 1,
+        },
         https: false,
+        http2: true,
         unix_path: None,
         access_log: None,
         lanes,
@@ -317,6 +372,44 @@ mod tests {
     use super::*;
     use crate::request::test_ctx;
 
+    #[test]
+    fn shards_shutdown_joins_threads_that_observe_the_signal() {
+        let (shutdown_tx, rx) = tokio::sync::watch::channel(false);
+        let thread = std::thread::spawn(move || {
+            while !*rx.borrow() {
+                std::thread::sleep(Duration::from_millis(1));
+            }
+        });
+
+        let start = Instant::now();
+        RuntimeHandle::Shards {
+            shutdown_tx,
+            threads: vec![thread],
+        }
+        .shutdown(Duration::from_secs(5));
+
+        // The thread exits on the signal, so the join comes nowhere near
+        // the deadline.
+        assert!(start.elapsed() < Duration::from_secs(1));
+    }
+
+    #[test]
+    fn shards_shutdown_abandons_a_wedged_thread_at_the_deadline() {
+        let (shutdown_tx, _rx) = tokio::sync::watch::channel(false);
+        let wedged = std::thread::spawn(|| std::thread::sleep(Duration::from_secs(30)));
+
+        let start = Instant::now();
+        RuntimeHandle::Shards {
+            shutdown_tx,
+            threads: vec![wedged],
+        }
+        .shutdown(Duration::from_millis(50));
+
+        let elapsed = start.elapsed();
+        assert!(elapsed >= Duration::from_millis(50));
+        assert!(elapsed < Duration::from_secs(5));
+    }
+
     #[test]
     fn worker_registration_hands_out_sequential_slot_ids() {
         let server = test_server(false, 4);
@@ -420,9 +513,9 @@ mod tests {
     #[test]
     fn queue_histogram_buckets_by_wait() {
         let h = QueueHistogram::new();
-        h.record(400);        // <= 500 -> bucket 0
-        h.record(500);        // == 500 -> bucket 0 (inclusive)
-        h.record(600);        // (500, 1000] -> bucket 1
+        h.record(400); // <= 500 -> bucket 0
+        h.record(500); // == 500 -> bucket 0 (inclusive)
+        h.record(600); // (500, 1000] -> bucket 1
         h.record(20_000_000); // > last bound -> overflow
         let s = h.snapshot();
         assert_eq!(s.buckets[0], 2);