itsi-server 0.2.15 → 0.2.17

data/ext/itsi_server/src/server/middleware_stack/middlewares/static_response.rs

@@ -1,15 +1,13 @@
 use std::sync::OnceLock;
 
 use super::{FromValue, MiddlewareLayer};
-use crate::server::http_message_types::{HttpRequest, HttpResponse};
+use crate::server::http_message_types::{HttpBody, HttpRequest, HttpResponse};
 use crate::services::itsi_http_service::HttpRequestContext;
 use async_trait::async_trait;
 use bytes::Bytes;
 use derive_more::Debug;
 use either::Either;
 use http::{HeaderMap, HeaderName, HeaderValue, Response, StatusCode};
-use http_body_util::combinators::BoxBody;
-use http_body_util::Full;
 use itsi_error::ItsiError;
 use magnus::error::Result;
 use serde::Deserialize;
@@ -21,6 +19,10 @@ pub struct StaticResponse {
     body: Vec<u8>,
     #[serde(skip)]
     header_map: OnceLock<HeaderMap>,
+    #[serde(skip)]
+    body_bytes: OnceLock<Bytes>,
+    #[serde(skip)]
+    status_code: OnceLock<StatusCode>,
 }
 
 #[async_trait]
@@ -35,6 +37,12 @@ impl MiddlewareLayer for StaticResponse {
         self.header_map
             .set(header_map)
             .map_err(|_| ItsiError::new("Failed to set headers"))?;
+        self.body_bytes
+            .set(Bytes::from(self.body.clone()))
+            .map_err(|_| ItsiError::new("Failed to set body bytes"))?;
+        self.status_code
+            .set(StatusCode::from_u16(self.code).unwrap_or(StatusCode::OK))
+            .map_err(|_| ItsiError::new("Failed to set status code"))?;
         Ok(())
     }
 
@@ -43,9 +51,8 @@ impl MiddlewareLayer for StaticResponse {
         _req: HttpRequest,
         _context: &mut HttpRequestContext,
     ) -> Result<Either<HttpRequest, HttpResponse>> {
-        let mut resp = Response::new(BoxBody::new(Full::new(Bytes::from(self.body.clone()))));
-        let status = StatusCode::from_u16(self.code).unwrap_or(StatusCode::OK);
-        *resp.status_mut() = status;
+        let mut resp = Response::new(HttpBody::full(self.body_bytes.get().unwrap().clone()));
+        *resp.status_mut() = *self.status_code.get().unwrap();
         *resp.headers_mut() = self.header_map.get().unwrap().clone();
 
         Ok(Either::Right(resp))

data/ext/itsi_server/src/server/mod.rs

@@ -1,5 +1,6 @@
 pub mod binds;
 pub mod byte_frame;
+pub mod frame_stream;
 pub mod http_message_types;
 pub mod io_stream;
 pub mod lifecycle_event;

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

@@ -42,7 +42,8 @@ impl Default for ProcessWorker {
     }
 }
 
-static CORE_IDS: LazyLock<Vec<CoreId>> = LazyLock::new(|| core_affinity::get_core_ids().unwrap());
+pub 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))]
@@ -78,7 +79,7 @@ impl ProcessWorker {
                 ) {
                     error!("Failed to set process group ID: {}", e);
                 }
-                match SingleMode::new(cluster_template.server_config.clone()) {
+                match SingleMode::new(cluster_template.server_config.clone(), self.worker_id) {
                     Ok(single_mode) => {
                         if cluster_template
                             .server_config
@@ -87,7 +88,7 @@ impl ProcessWorker {
                             .pin_worker_cores
                         {
                             core_affinity::set_for_current(
-                                CORE_IDS[self.worker_id % CORE_IDS.len()],
+                                CORE_IDS[(2 * self.worker_id) % CORE_IDS.len()],
                             );
                         }
                         Arc::new(single_mode).run().ok();
@@ -165,7 +166,7 @@ impl ProcessWorker {
     }
 
     pub(crate) fn boot_if_dead(&self, cluster_template: Arc<ClusterMode>) -> bool {
-        if !self.is_alive() {
+        if !self.is_alive() && self.child_pid.lock().is_some() {
             if self.just_started() {
                 error!(
                     "Worker in crash loop {:?}. Refusing to restart",
@@ -201,7 +202,6 @@ 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/ext/itsi_server/src/server/serve_strategy/acceptor.rs

@@ -0,0 +1,100 @@
+use hyper_util::rt::TokioIo;
+use std::{ops::Deref, pin::Pin, sync::Arc, time::Duration};
+use tokio::task::JoinSet;
+use tracing::debug;
+
+use crate::{
+    ruby_types::itsi_server::itsi_server_config::ServerParams,
+    server::{binds::listener::ListenerInfo, io_stream::IoStream, request_job::RequestJob},
+    services::itsi_http_service::{ItsiHttpService, ItsiHttpServiceInner},
+};
+
+use super::single_mode::{RunningPhase, SingleMode};
+
+pub struct Acceptor {
+    pub acceptor_args: Arc<AcceptorArgs>,
+    pub join_set: JoinSet<()>,
+}
+
+impl Deref for Acceptor {
+    type Target = Arc<AcceptorArgs>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.acceptor_args
+    }
+}
+
+pub struct AcceptorArgs {
+    pub strategy: Arc<SingleMode>,
+    pub listener_info: ListenerInfo,
+    pub shutdown_receiver: tokio::sync::watch::Receiver<RunningPhase>,
+    pub job_sender: async_channel::Sender<RequestJob>,
+    pub nonblocking_sender: async_channel::Sender<RequestJob>,
+    pub server_params: Arc<ServerParams>,
+}
+
+impl Acceptor {
+    pub(crate) async fn serve_connection(&mut self, stream: IoStream) {
+        let addr = stream.addr();
+        let io: TokioIo<Pin<Box<IoStream>>> = TokioIo::new(Box::pin(stream));
+        let mut shutdown_channel = self.shutdown_receiver.clone();
+        let acceptor_args = self.acceptor_args.clone();
+        let service = ItsiHttpService {
+            inner: Arc::new(ItsiHttpServiceInner {
+                acceptor_args: acceptor_args.clone(),
+                addr,
+            }),
+        };
+
+        self.join_set.spawn(async move {
+            let executor = &acceptor_args.strategy.executor;
+            let svc = hyper::service::service_fn(move |req| {
+                let service = service.clone();
+                async move { service.handle_request(req).await }
+            });
+
+            let mut serve = Box::pin(executor.serve_connection_with_upgrades(io, svc));
+
+            tokio::select! {
+                // Await the connection finishing naturally.
+                res = &mut serve => {
+                    match res {
+                        Ok(()) => {
+                            debug!("Connection closed normally");
+                        },
+                        Err(res) => {
+                            debug!("Connection closed abruptly: {:?}", res);
+                        }
+                    }
+                },
+                // A lifecycle event triggers shutdown.
+                _ = shutdown_channel.changed() => {
+                    // Initiate graceful shutdown.
+                    serve.as_mut().graceful_shutdown();
+
+                    // Now await the connection to finish shutting down.
+                    if let Err(e) = serve.await {
+                        debug!("Connection shutdown error: {:?}", e);
+                    }
+                }
+            }
+        });
+    }
+
+    pub async fn join(&mut self) {
+        // Join all acceptor tasks with timeout
+
+        let deadline = tokio::time::Instant::now()
+            + Duration::from_secs_f64(self.server_params.shutdown_timeout);
+        let sleep_until = tokio::time::sleep_until(deadline);
+        tokio::select! {
+            _ = async {
+                while (self.join_set.join_next().await).is_some() {}
+            } => {},
+            _ = sleep_until => {
+                self.join_set.abort_all();
+                debug!("Shutdown timeout reached; abandoning remaining acceptor tasks.");
+            }
+        }
+    }
+}

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

@@ -1,5 +1,5 @@
 use crate::ruby_types::itsi_server::itsi_server_config::ItsiServerConfig;
-use crate::server::signal::SIGNAL_HANDLER_CHANNEL;
+use crate::server::signal::{subscribe_runtime_to_signals, unsubscribe_runtime};
 use crate::server::{lifecycle_event::LifecycleEvent, process_worker::ProcessWorker};
 use itsi_error::{ItsiError, Result};
 use itsi_rb_helpers::{call_with_gvl, call_without_gvl, create_ruby_thread};
@@ -7,31 +7,32 @@ use itsi_tracing::{error, info, warn};
 use magnus::Value;
 use nix::{libc::exit, unistd::Pid};
 
+use std::sync::atomic::{AtomicBool, Ordering};
 use std::{
-    sync::{atomic::AtomicUsize, Arc},
+    sync::Arc,
     time::{Duration, Instant},
 };
 use tokio::{
     runtime::{Builder as RuntimeBuilder, Runtime},
-    sync::{broadcast, watch, Mutex},
+    sync::{watch, Mutex},
     time::{self, sleep},
 };
 use tracing::{debug, instrument};
 pub(crate) struct ClusterMode {
     pub server_config: Arc<ItsiServerConfig>,
     pub process_workers: parking_lot::Mutex<Vec<ProcessWorker>>,
-    pub lifecycle_channel: broadcast::Sender<LifecycleEvent>,
 }
 
-static WORKER_ID: AtomicUsize = AtomicUsize::new(0);
 static CHILD_SIGNAL_SENDER: parking_lot::Mutex<Option<watch::Sender<()>>> =
     parking_lot::Mutex::new(None);
 
+static RELOAD_IN_PROGRESS: AtomicBool = AtomicBool::new(false);
+
 impl ClusterMode {
     pub fn new(server_config: Arc<ItsiServerConfig>) -> Self {
         let process_workers = (0..server_config.server_params.read().workers)
-            .map(|_| ProcessWorker {
-                worker_id: WORKER_ID.fetch_add(1, std::sync::atomic::Ordering::Relaxed),
+            .map(|id| ProcessWorker {
+                worker_id: id as usize,
                 ..Default::default()
             })
             .collect();
@@ -39,7 +40,6 @@ impl ClusterMode {
         Self {
             server_config,
             process_workers: parking_lot::Mutex::new(process_workers),
-            lifecycle_channel: SIGNAL_HANDLER_CHANNEL.0.clone(),
         }
     }
 
@@ -60,6 +60,26 @@ impl ClusterMode {
         }
     }
 
+    fn next_worker_id(&self) -> usize {
+        let mut ids: Vec<usize> = self
+            .process_workers
+            .lock()
+            .iter()
+            .map(|w| w.worker_id)
+            .collect();
+        self.next_available_id_in(&mut ids)
+    }
+
+    fn next_available_id_in(&self, list: &mut [usize]) -> usize {
+        list.sort_unstable();
+        for (expected, &id) in list.iter().enumerate() {
+            if id != expected {
+                return expected;
+            }
+        }
+        list.len()
+    }
+
     #[allow(clippy::await_holding_lock)]
     pub async fn handle_lifecycle_event(
         self: Arc<Self>,
@@ -97,40 +117,56 @@ impl ClusterMode {
                     self.shutdown().await.ok();
                     self.server_config.reload_exec()?;
                 }
-                let mut workers_to_load = self.server_config.server_params.read().workers;
-                let mut next_workers = Vec::new();
-                for worker in self.process_workers.lock().drain(..) {
-                    if workers_to_load == 0 {
-                        worker.graceful_shutdown(self.clone()).await
-                    } else {
-                        workers_to_load -= 1;
-                        worker.reboot(self.clone()).await?;
-                        next_workers.push(worker);
-                    }
+
+                if RELOAD_IN_PROGRESS
+                    .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
+                    .is_err()
+                {
+                    warn!("Reload already in progress, ignoring request");
+                    return Ok(());
                 }
-                self.process_workers.lock().extend(next_workers);
-                while workers_to_load > 0 {
-                    let mut workers = self.process_workers.lock();
+                let workers_to_load = self.server_config.server_params.read().workers;
+                let mut next_workers = Vec::new();
+                let mut old_workers = self.process_workers.lock().drain(..).collect::<Vec<_>>();
+
+                // Spawn new workers
+                for i in 0..workers_to_load {
                     let worker = ProcessWorker {
-                        worker_id: WORKER_ID.fetch_add(1, std::sync::atomic::Ordering::Relaxed),
+                        worker_id: i as usize,
                         ..Default::default()
                     };
                     let worker_clone = worker.clone();
                     let self_clone = self.clone();
-                    create_ruby_thread(move || {
-                        call_without_gvl(move || {
-                            worker_clone.boot(self_clone).ok();
-                        })
+
+                    call_with_gvl(|_| {
+                        create_ruby_thread(move || {
+                            call_without_gvl(move || match worker_clone.boot(self_clone) {
+                                Err(err) => error!("Worker boot failed {:?}", err),
+                                _ => {}
+                            })
+                        });
                     });
-                    workers.push(worker);
-                    workers_to_load -= 1
+
+                    next_workers.push(worker);
+
+                    if let Some(old) = old_workers.pop() {
+                        old.graceful_shutdown(self.clone()).await;
+                    }
                 }
+
+                for worker in old_workers {
+                    worker.graceful_shutdown(self.clone()).await;
+                }
+
+                self.process_workers.lock().extend(next_workers);
+                RELOAD_IN_PROGRESS.store(false, Ordering::SeqCst);
+
                 Ok(())
             }
             LifecycleEvent::IncreaseWorkers => {
                 let mut workers = self.process_workers.lock();
                 let worker = ProcessWorker {
-                    worker_id: WORKER_ID.fetch_add(1, std::sync::atomic::Ordering::Relaxed),
+                    worker_id: self.next_worker_id(),
                     ..Default::default()
                 };
                 let worker_clone = worker.clone();
@@ -171,6 +207,10 @@ impl ClusterMode {
                 unsafe { exit(0) };
             }
             LifecycleEvent::ChildTerminated => {
+                if RELOAD_IN_PROGRESS.load(Ordering::SeqCst) {
+                    warn!("Reload already in progress, ignoring child signal");
+                    return Ok(());
+                }
                 CHILD_SIGNAL_SENDER.lock().as_ref().inspect(|i| {
                     i.send(()).ok();
                 });
@@ -275,18 +315,29 @@ impl ClusterMode {
     pub fn run(self: Arc<Self>) -> Result<()> {
         info!("Starting in Cluster mode");
         self.invoke_hook("before_fork");
+
         self.process_workers
             .lock()
             .iter()
             .try_for_each(|worker| worker.boot(Arc::clone(&self)))?;
 
+        if cfg!(target_os = "linux") {
+            self.server_config
+                .server_params
+                .write()
+                .listeners
+                .lock()
+                .drain(..);
+        };
+
         let (sender, mut receiver) = watch::channel(());
         *CHILD_SIGNAL_SENDER.lock() = Some(sender);
 
-        let mut lifecycle_rx = self.lifecycle_channel.subscribe();
         let self_ref = self.clone();
 
         self.build_runtime().block_on(async {
+          let mut lifecycle_rx = subscribe_runtime_to_signals();
+
           let self_ref = self_ref.clone();
           let memory_check_duration = if self_ref.server_config.server_params.read().worker_memory_limit.is_some(){
             time::Duration::from_secs(15)
@@ -338,11 +389,16 @@ impl ClusterMode {
                   }
 
                 },
-                Err(e) => error!("Error receiving lifecycle_event: {:?}", e),
+                Err(e) => {
+                  debug!("Lifecycle channel closed: {:?}, exiting cluster monitor loop", e);
+                  break
+                },
               }
             }
           }
         });
+
+        unsubscribe_runtime();
         self.server_config
             .server_params
             .write()

data/ext/itsi_server/src/server/serve_strategy/mod.rs

@@ -4,6 +4,7 @@ use cluster_mode::ClusterMode;
 use itsi_error::Result;
 use single_mode::SingleMode;
 
+pub mod acceptor;
 pub mod cluster_mode;
 pub mod single_mode;