itsi 0.1.14 → 0.1.18

data/crates/itsi_server/src/server/process_worker.rs

@@ -1,4 +1,5 @@
 use super::serve_strategy::{cluster_mode::ClusterMode, single_mode::SingleMode};
+use core_affinity::CoreId;
 use itsi_error::{ItsiError, Result};
 use itsi_rb_helpers::{call_with_gvl, call_without_gvl, create_ruby_thread, fork};
 use itsi_tracing::error;
@@ -16,7 +17,7 @@ use nix::{
 use parking_lot::Mutex;
 use std::{
     process::{self, exit},
-    sync::Arc,
+    sync::{Arc, LazyLock},
     time::{Duration, Instant},
 };
 use sysinfo::System;
@@ -41,6 +42,8 @@ impl Default for ProcessWorker {
     }
 }
 
+static CORE_IDS: LazyLock<Vec<CoreId>> = LazyLock::new(|| core_affinity::get_core_ids().unwrap());
+
 impl ProcessWorker {
     #[instrument(skip(self, cluster_template), fields(self.worker_id = %self.worker_id))]
     pub(crate) fn boot(&self, cluster_template: Arc<ClusterMode>) -> Result<()> {
@@ -77,6 +80,16 @@ impl ProcessWorker {
                 }
                 match SingleMode::new(cluster_template.server_config.clone()) {
                     Ok(single_mode) => {
+                        if cluster_template
+                            .server_config
+                            .server_params
+                            .read()
+                            .pin_worker_cores
+                        {
+                            core_affinity::set_for_current(
+                                CORE_IDS[self.worker_id % CORE_IDS.len()],
+                            );
+                        }
                         Arc::new(single_mode).run().ok();
                     }
                     Err(e) => {
@@ -174,8 +187,12 @@ impl ProcessWorker {
     pub(crate) fn request_shutdown(&self) {
         let child_pid = *self.child_pid.lock();
         if let Some(pid) = child_pid {
-            if let Err(e) = kill(pid, SIGTERM) {
-                error!("Failed to send SIGTERM to process {}: {}", pid, e);
+            if self.is_alive() {
+                if let Err(e) = kill(pid, SIGTERM) {
+                    error!("Failed to send SIGTERM to process {}: {}", pid, e);
+                }
+            } else {
+                error!("Trying to shutdown a dead process");
             }
         }
     }
@@ -184,6 +201,7 @@ impl ProcessWorker {
         let child_pid = *self.child_pid.lock();
         if let Some(pid) = child_pid {
             if self.is_alive() {
+                info!("Worker still alive, sending SIGKILL {}", pid);
                 if let Err(e) = kill(pid, SIGKILL) {
                     error!("Failed to force kill process {}: {}", pid, e);
                 }

data/crates/itsi_server/src/server/request_job.rs

@@ -1,4 +1,4 @@
-use crate::ruby_types::{itsi_grpc_request::ItsiGrpcRequest, itsi_http_request::ItsiHttpRequest};
+use crate::ruby_types::{itsi_grpc_call::ItsiGrpcCall, itsi_http_request::ItsiHttpRequest};
 use itsi_rb_helpers::HeapValue;
 use magnus::block::Proc;
 use std::sync::Arc;
@@ -6,6 +6,6 @@ use std::sync::Arc;
 #[derive(Debug)]
 pub enum RequestJob {
     ProcessHttpRequest(ItsiHttpRequest, Arc<HeapValue<Proc>>),
-    ProcessGrpcRequest(ItsiGrpcRequest, Arc<HeapValue<Proc>>),
+    ProcessGrpcRequest(ItsiGrpcCall, Arc<HeapValue<Proc>>),
     Shutdown,
 }

data/crates/itsi_server/src/server/serve_strategy/cluster_mode.rs

@@ -202,7 +202,7 @@ impl ClusterMode {
             }
         }
 
-        Err(ItsiError::Break())
+        Err(ItsiError::Break)
     }
 
     pub async fn print_info(self: Arc<Self>) -> Result<()> {
@@ -320,7 +320,7 @@ impl ClusterMode {
                 Ok(lifecycle_event) => {
                   if let Err(e) = self_ref.clone().handle_lifecycle_event(lifecycle_event).await{
                     match e {
-                      ItsiError::Break() => break,
+                      ItsiError::Break => break,
                       _ => error!("Error in handle_lifecycle_event {:?}", e)
                     }
                   }
@@ -331,7 +331,12 @@ impl ClusterMode {
             }
           }
         });
-
+        self.server_config
+            .server_params
+            .write()
+            .listeners
+            .lock()
+            .drain(..);
         Ok(())
     }
 }

data/crates/itsi_server/src/server/serve_strategy/single_mode.rs

@@ -1,14 +1,14 @@
 use crate::{
     ruby_types::itsi_server::itsi_server_config::ItsiServerConfig,
     server::{
+        binds::listener::ListenerInfo,
         io_stream::IoStream,
-        itsi_service::{IstiServiceInner, ItsiService},
         lifecycle_event::LifecycleEvent,
-        listener::ListenerInfo,
         request_job::RequestJob,
-        signal::SIGNAL_HANDLER_CHANNEL,
+        signal::{SHUTDOWN_REQUESTED, SIGNAL_HANDLER_CHANNEL},
         thread_worker::{build_thread_workers, ThreadWorker},
     },
+    services::itsi_http_service::{ItsiHttpService, ItsiHttpServiceInner},
 };
 use hyper_util::{
     rt::{TokioExecutor, TokioIo, TokioTimer},
@@ -24,7 +24,6 @@ use nix::unistd::Pid;
 use parking_lot::RwLock;
 use std::{
     collections::HashMap,
-    num::NonZeroU8,
     pin::Pin,
     sync::{
         atomic::{AtomicBool, Ordering},
@@ -102,10 +101,6 @@ impl SingleMode {
             self.server_config.server_params.read().binds
         );
 
-        println!(
-            "    ─ script_name: {:?}",
-            self.server_config.server_params.read().script_name
-        );
         println!(
             "    ─ streaming body: {:?}",
             self.server_config.server_params.read().streamable_body
@@ -222,25 +217,24 @@ impl SingleMode {
         let mut listener_task_set = JoinSet::new();
         let runtime = self.build_runtime();
 
-        let (thread_workers, job_sender) = build_thread_workers(
-            self.server_config.server_params.read().clone(),
-            Pid::this(),
-            NonZeroU8::try_from(self.server_config.server_params.read().threads).unwrap(),
-        )
-        .inspect_err(|e| {
-            if let Some(err_val) = e.value() {
-                print_rb_backtrace(err_val);
-            }
-        })?;
+        let (thread_workers, job_sender, nonblocking_sender) =
+            build_thread_workers(self.server_config.server_params.read().clone(), Pid::this())
+                .inspect_err(|e| {
+                    if let Some(err_val) = e.value() {
+                        print_rb_backtrace(err_val);
+                    }
+                })?;
 
         info!(
-            pid = format!("{}", Pid::this()),
             threads = thread_workers.len(),
             binds = format!("{:?}", self.server_config.server_params.read().binds)
         );
 
         let (shutdown_sender, _) = watch::channel(RunningPhase::Running);
         let thread = self.clone().start_monitors(thread_workers.clone());
+        if SHUTDOWN_REQUESTED.load(Ordering::SeqCst) {
+            return Ok(());
+        }
         runtime.block_on(
           async  {
               let server_params = self.server_config.server_params.read().clone();
@@ -254,11 +248,13 @@ impl SingleMode {
 
               for listener in tokio_listeners.iter() {
                   let mut lifecycle_rx = self.lifecycle_channel.subscribe();
+
                   let listener_info = Arc::new(listener.listener_info());
                   let self_ref = self.clone();
                   let listener = listener.clone();
                   let shutdown_sender = shutdown_sender.clone();
                   let job_sender = job_sender.clone();
+                  let nonblocking_sender = nonblocking_sender.clone();
                   let workers_clone = thread_workers.clone();
                   let listener_clone = listener.clone();
                   let mut shutdown_receiver = shutdown_sender.subscribe();
@@ -278,8 +274,9 @@ impl SingleMode {
                                 let listener_info = listener_info.clone();
                                 let shutdown_receiver = shutdown_receiver.clone();
                                 let job_sender = job_sender.clone();
+                                let nonblocking_sender = nonblocking_sender.clone();
                                 acceptor_task_set.spawn(async move {
-                                  strategy.serve_connection(accept_result, job_sender, listener_info, shutdown_receiver).await;
+                                  strategy.serve_connection(accept_result, job_sender, nonblocking_sender, listener_info, shutdown_receiver).await;
                                 });
                               },
                               Err(e) => debug!("Listener.accept failed {:?}", e),
@@ -290,11 +287,10 @@ impl SingleMode {
                             lifecycle_event = lifecycle_rx.recv() => match lifecycle_event{
                               Ok(LifecycleEvent::Shutdown) => {
                                 shutdown_sender.send(RunningPhase::ShutdownPending).unwrap();
-                                // Tell any in-progress connections to stop accepting new requests
                                 tokio::time::sleep(Duration::from_millis(25)).await;
-                                // Tell workers to stop processing requests once they've flushed their buffers.
                                 for _i in 0..workers_clone.len() {
                                   job_sender.send(RequestJob::Shutdown).await.unwrap();
+                                  nonblocking_sender.send(RequestJob::Shutdown).await.unwrap();
                                 }
                                 break;
                               },
@@ -310,6 +306,9 @@ impl SingleMode {
 
               while let Some(_res) = listener_task_set.join_next().await {}
 
+              // Explicitly drop all listeners to ensure file descriptors are released
+              drop(tokio_listeners);
+
               Ok::<(), ItsiError>(())
           })?;
 
@@ -343,6 +342,7 @@ impl SingleMode {
         &self,
         stream: IoStream,
         job_sender: async_channel::Sender<RequestJob>,
+        nonblocking_sender: async_channel::Sender<RequestJob>,
         listener: Arc<ListenerInfo>,
         shutdown_channel: watch::Receiver<RunningPhase>,
     ) {
@@ -354,9 +354,10 @@ impl SingleMode {
         let mut binding = executor.http1();
         let shutdown_channel = shutdown_channel_clone.clone();
 
-        let service = ItsiService {
-            inner: Arc::new(IstiServiceInner {
+        let service = ItsiHttpService {
+            inner: Arc::new(ItsiHttpServiceInner {
                 sender: job_sender.clone(),
+                nonblocking_sender: nonblocking_sender.clone(),
                 server_params: self.server_config.server_params.read().clone(),
                 listener,
                 addr: addr.to_string(),
@@ -365,7 +366,7 @@ impl SingleMode {
         };
         let mut serve = Box::pin(
             binding
-                .timer(TokioTimer::new()) // your existing timer
+                .timer(TokioTimer::new())
                 .header_read_timeout(Duration::from_secs(1))
                 .serve_connection_with_upgrades(io, service),
         );
@@ -386,7 +387,6 @@ impl SingleMode {
             // A lifecycle event triggers shutdown.
             _ = shutdown_channel_clone.changed() => {
                 // Initiate graceful shutdown.
-                info!("Starting graceful shutdown");
                 serve.as_mut().graceful_shutdown();
 
                 // Now await the connection to finish shutting down.

data/crates/itsi_server/src/server/signal.rs

@@ -1,71 +1,54 @@
-use std::sync::{atomic::AtomicI8, LazyLock};
+use std::sync::{
+    atomic::{AtomicBool, AtomicI8},
+    LazyLock,
+};
 
 use nix::libc::{self, sighandler_t};
 use tokio::sync::{self, broadcast};
 
 use super::lifecycle_event::LifecycleEvent;
 
+pub static SIGINT_COUNT: AtomicI8 = AtomicI8::new(0);
+pub static SHUTDOWN_REQUESTED: AtomicBool = AtomicBool::new(false);
 pub static SIGNAL_HANDLER_CHANNEL: LazyLock<(
     broadcast::Sender<LifecycleEvent>,
     broadcast::Receiver<LifecycleEvent>,
 )> = LazyLock::new(|| sync::broadcast::channel(5));
 
-pub fn send_shutdown_event() {
-    SIGNAL_HANDLER_CHANNEL.0.send(LifecycleEvent::Shutdown).ok();
+pub fn send_lifecycle_event(event: LifecycleEvent) {
+    SIGNAL_HANDLER_CHANNEL.0.send(event).ok();
 }
 
-pub static SIGINT_COUNT: AtomicI8 = AtomicI8::new(0);
 fn receive_signal(signum: i32, _: sighandler_t) {
     SIGINT_COUNT.fetch_add(-1, std::sync::atomic::Ordering::SeqCst);
-    match signum {
+    let event = match signum {
         libc::SIGTERM | libc::SIGINT => {
+            SHUTDOWN_REQUESTED.store(true, std::sync::atomic::Ordering::SeqCst);
             SIGINT_COUNT.fetch_add(2, std::sync::atomic::Ordering::SeqCst);
             if SIGINT_COUNT.load(std::sync::atomic::Ordering::SeqCst) < 2 {
-                SIGNAL_HANDLER_CHANNEL.0.send(LifecycleEvent::Shutdown).ok();
+                Some(LifecycleEvent::Shutdown)
             } else {
                 // Not messing about. Force shutdown.
-                SIGNAL_HANDLER_CHANNEL
-                    .0
-                    .send(LifecycleEvent::ForceShutdown)
-                    .ok();
+                Some(LifecycleEvent::ForceShutdown)
             }
         }
-        libc::SIGUSR2 => {
-            SIGNAL_HANDLER_CHANNEL
-                .0
-                .send(LifecycleEvent::PrintInfo)
-                .ok();
-        }
-        libc::SIGUSR1 => {
-            SIGNAL_HANDLER_CHANNEL.0.send(LifecycleEvent::Restart).ok();
-        }
-        libc::SIGHUP => {
-            SIGNAL_HANDLER_CHANNEL.0.send(LifecycleEvent::Reload).ok();
-        }
-        libc::SIGTTIN => {
-            SIGNAL_HANDLER_CHANNEL
-                .0
-                .send(LifecycleEvent::IncreaseWorkers)
-                .ok();
-        }
-        libc::SIGTTOU => {
-            SIGNAL_HANDLER_CHANNEL
-                .0
-                .send(LifecycleEvent::DecreaseWorkers)
-                .ok();
-        }
-        libc::SIGCHLD => {
-            SIGNAL_HANDLER_CHANNEL
-                .0
-                .send(LifecycleEvent::ChildTerminated)
-                .ok();
-        }
-        _ => {}
+        libc::SIGUSR2 => Some(LifecycleEvent::PrintInfo),
+        libc::SIGUSR1 => Some(LifecycleEvent::Restart),
+        libc::SIGHUP => Some(LifecycleEvent::Reload),
+        libc::SIGTTIN => Some(LifecycleEvent::IncreaseWorkers),
+        libc::SIGTTOU => Some(LifecycleEvent::DecreaseWorkers),
+        libc::SIGCHLD => Some(LifecycleEvent::ChildTerminated),
+        _ => None,
+    };
+
+    if let Some(event) = event {
+        send_lifecycle_event(event);
     }
 }
 
 pub fn reset_signal_handlers() -> bool {
     SIGINT_COUNT.store(0, std::sync::atomic::Ordering::SeqCst);
+    SHUTDOWN_REQUESTED.store(false, std::sync::atomic::Ordering::SeqCst);
     unsafe {
         libc::signal(libc::SIGTERM, receive_signal as usize);
         libc::signal(libc::SIGINT, receive_signal as usize);

data/crates/itsi_server/src/server/size_limited_incoming.rs

@@ -0,0 +1,101 @@
+use bytes::Buf;
+use hyper::body::Body;
+use hyper::body::Frame;
+use hyper::body::SizeHint;
+use std::error::Error;
+use std::fmt;
+use std::ops::Deref;
+use std::pin::Pin;
+use std::sync::atomic::AtomicUsize;
+use std::sync::atomic::Ordering;
+use std::task::Context;
+use std::task::Poll;
+
+/// Custom error to indicate that the maximum body size was exceeded.
+#[derive(Debug)]
+pub struct MaxBodySizeReached;
+impl fmt::Display for MaxBodySizeReached {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "Maximum body size reached")
+    }
+}
+
+impl Error for MaxBodySizeReached {}
+
+#[derive(Debug)]
+pub struct SizeLimitedIncoming<B> {
+    pub inner: B,
+    pub limit: AtomicUsize,
+    current: usize,
+}
+
+impl<B> Deref for SizeLimitedIncoming<B> {
+    type Target = B;
+
+    fn deref(&self) -> &Self::Target {
+        &self.inner
+    }
+}
+
+impl<B> SizeLimitedIncoming<B> {
+    pub fn new(inner: B) -> Self {
+        Self {
+            inner,
+            limit: AtomicUsize::new(usize::MAX),
+            current: 0,
+        }
+    }
+}
+
+impl<B> Body for SizeLimitedIncoming<B>
+where
+    B: Body + Unpin,
+    B::Data: Buf,
+    // Ensure that the inner error converts into our boxed error type.
+    B::Error: Into<Box<dyn Error + Send + Sync>>,
+{
+    type Data = B::Data;
+    type Error = Box<dyn Error + Send + Sync>;
+
+    fn poll_frame(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
+        // Pin the inner body.
+        let inner = Pin::new(&mut self.inner);
+        match inner.poll_frame(cx) {
+            Poll::Ready(Some(Ok(frame))) => {
+                // Use public methods since we cannot match on the private enum.
+                if frame.is_data() {
+                    match frame.into_data() {
+                        Ok(data) => {
+                            let len = data.remaining();
+                            self.current += len;
+                            if self.current > self.limit.load(Ordering::Relaxed) {
+                                Poll::Ready(Some(Err(Box::new(MaxBodySizeReached))))
+                            } else {
+                                Poll::Ready(Some(Ok(Frame::data(data))))
+                            }
+                        }
+                        // Should not occur if is_data() was true, but pass through if it does.
+                        Err(frame) => Poll::Ready(Some(Ok(frame))),
+                    }
+                } else {
+                    // For non-data frames (e.g. trailers), just pass them along.
+                    Poll::Ready(Some(Ok(frame)))
+                }
+            }
+            Poll::Ready(Some(Err(e))) => Poll::Ready(Some(Err(e.into()))),
+            Poll::Ready(None) => Poll::Ready(None),
+            Poll::Pending => Poll::Pending,
+        }
+    }
+
+    fn is_end_stream(&self) -> bool {
+        self.inner.is_end_stream()
+    }
+
+    fn size_hint(&self) -> SizeHint {
+        self.inner.size_hint()
+    }
+}

data/crates/itsi_server/src/server/thread_worker.rs

@@ -11,7 +11,6 @@ use magnus::{
 use nix::unistd::Pid;
 use parking_lot::{Mutex, RwLock};
 use std::{
-    num::NonZeroU8,
     ops::Deref,
     sync::{
         atomic::{AtomicBool, AtomicU64, Ordering},
@@ -24,7 +23,7 @@ use tokio::{runtime::Builder as RuntimeBuilder, sync::watch};
 use tracing::instrument;
 
 use crate::ruby_types::{
-    itsi_grpc_request::ItsiGrpcRequest, itsi_http_request::ItsiHttpRequest,
+    itsi_grpc_call::ItsiGrpcCall, itsi_http_request::ItsiHttpRequest,
     itsi_server::itsi_server_config::ServerParams, ITSI_SERVER,
 };
 
@@ -55,34 +54,66 @@ static CLASS_FIBER: Lazy<RClass> = Lazy::new(|ruby| {
 });
 
 pub struct TerminateWakerSignal(bool);
-type ThreadWorkerBuildResult = Result<(Arc<Vec<Arc<ThreadWorker>>>, Sender<RequestJob>)>;
+type ThreadWorkerBuildResult = Result<(
+    Arc<Vec<Arc<ThreadWorker>>>,
+    Sender<RequestJob>,
+    Sender<RequestJob>,
+)>;
 
-#[instrument(name = "boot", parent=None, skip(params, threads, pid))]
-pub fn build_thread_workers(
-    params: Arc<ServerParams>,
-    pid: Pid,
-    threads: NonZeroU8,
-) -> ThreadWorkerBuildResult {
-    let (sender, receiver) = async_channel::bounded((threads.get() as u16 * 30) as usize);
-    let receiver_ref = Arc::new(receiver);
-    let sender_ref = sender;
+#[instrument(name = "boot", parent=None, skip(params, pid))]
+pub fn build_thread_workers(params: Arc<ServerParams>, pid: Pid) -> ThreadWorkerBuildResult {
+    let blocking_thread_count = params.threads;
+    let nonblocking_thread_count = params.scheduler_threads;
+
+    let (blocking_sender, blocking_receiver) =
+        async_channel::bounded((blocking_thread_count as u16 * 30) as usize);
+    let blocking_receiver_ref = Arc::new(blocking_receiver);
+    let blocking_sender_ref = blocking_sender;
     let scheduler_class = load_scheduler_class(params.scheduler_class.clone())?;
+
+    let mut workers = (1..=blocking_thread_count)
+        .map(|id| {
+            ThreadWorker::new(
+                params.clone(),
+                id,
+                format!("{:?}#{:?}", pid, id),
+                blocking_receiver_ref.clone(),
+                blocking_sender_ref.clone(),
+                if nonblocking_thread_count.is_some() {
+                    None
+                } else {
+                    scheduler_class
+                },
+            )
+        })
+        .collect::<Result<Vec<_>>>()?;
+
+    let nonblocking_sender_ref = if let (Some(nonblocking_thread_count), Some(scheduler_class)) =
+        (nonblocking_thread_count, scheduler_class)
+    {
+        let (nonblocking_sender, nonblocking_receiver) =
+            async_channel::bounded((nonblocking_thread_count as u16 * 30) as usize);
+        let nonblocking_receiver_ref = Arc::new(nonblocking_receiver);
+        let nonblocking_sender_ref = nonblocking_sender.clone();
+        for id in 0..nonblocking_thread_count {
+            workers.push(ThreadWorker::new(
+                params.clone(),
+                id,
+                format!("{:?}#{:?}", pid, id),
+                nonblocking_receiver_ref.clone(),
+                nonblocking_sender_ref.clone(),
+                Some(scheduler_class),
+            )?)
+        }
+        nonblocking_sender
+    } else {
+        blocking_sender_ref.clone()
+    };
+
     Ok((
-        Arc::new(
-            (1..=u8::from(threads))
-                .map(|id| {
-                    ThreadWorker::new(
-                        params.clone(),
-                        id,
-                        format!("{:?}#{:?}", pid, id),
-                        receiver_ref.clone(),
-                        sender_ref.clone(),
-                        scheduler_class,
-                    )
-                })
-                .collect::<Result<Vec<_>>>()?,
-        ),
-        sender_ref,
+        Arc::new(workers),
+        blocking_sender_ref,
+        nonblocking_sender_ref,
     ))
 }
 
@@ -259,7 +290,7 @@ impl ThreadWorker {
                                     *ID_SCHEDULE,
                                     (app_proc.as_value(), request),
                                 ) {
-                                    ItsiGrpcRequest::internal_error(ruby, response, err)
+                                    ItsiGrpcCall::internal_error(ruby, response, err)
                                 }
                             }
                             RequestJob::Shutdown => return true,