itsi 0.1.19 → 0.2.2

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

@@ -2,9 +2,7 @@ use crate::ruby_types::itsi_server::itsi_server_config::ItsiServerConfig;
 use crate::server::signal::SIGNAL_HANDLER_CHANNEL;
 use crate::server::{lifecycle_event::LifecycleEvent, process_worker::ProcessWorker};
 use itsi_error::{ItsiError, Result};
-use itsi_rb_helpers::{
-    call_proc_and_log_errors, call_with_gvl, call_without_gvl, create_ruby_thread,
-};
+use itsi_rb_helpers::{call_with_gvl, call_without_gvl, create_ruby_thread};
 use itsi_tracing::{error, info, warn};
 use magnus::Value;
 use nix::{libc::exit, unistd::Pid};
@@ -56,6 +54,12 @@ impl ClusterMode {
             .expect("Failed to build Tokio runtime")
     }
 
+    pub fn invoke_hook(&self, hook_name: &str) {
+        if let Some(hook) = self.server_config.server_params.read().hooks.get(hook_name) {
+            call_with_gvl(|_| hook.call::<_, Value>(()).ok());
+        }
+    }
+
     #[allow(clippy::await_holding_lock)]
     pub async fn handle_lifecycle_event(
         self: Arc<Self>,
@@ -70,16 +74,23 @@ impl ClusterMode {
             LifecycleEvent::Shutdown => {
                 self.server_config.stop_watcher()?;
                 self.shutdown().await?;
+                self.invoke_hook("before_shutdown");
                 Ok(())
             }
             LifecycleEvent::Restart => {
-                self.server_config.dup_fds()?;
-                self.shutdown().await.ok();
-                info!("Shutdown complete. Calling reload exec");
-                self.server_config.reload_exec()?;
+                if self.server_config.check_config().await {
+                    self.invoke_hook("before_restart");
+                    self.server_config.dup_fds()?;
+                    self.shutdown().await.ok();
+                    info!("Shutdown complete. Calling reload exec");
+                    self.server_config.reload_exec()?;
+                }
                 Ok(())
             }
             LifecycleEvent::Reload => {
+                if !self.server_config.check_config().await {
+                    return Ok(());
+                }
                 let should_reexec = self.server_config.clone().reload(true)?;
                 if should_reexec {
                     self.server_config.dup_fds()?;
@@ -263,15 +274,7 @@ impl ClusterMode {
     #[instrument(skip(self), fields(mode = "cluster", pid=format!("{:?}", Pid::this())))]
     pub fn run(self: Arc<Self>) -> Result<()> {
         info!("Starting in Cluster mode");
-        if let Some(proc) = self
-            .server_config
-            .server_params
-            .read()
-            .hooks
-            .get("before_fork")
-        {
-            call_with_gvl(|_| call_proc_and_log_errors(proc.clone()))
-        }
+        self.invoke_hook("before_fork");
         self.process_workers
             .lock()
             .iter()
@@ -285,7 +288,16 @@ impl ClusterMode {
 
         self.build_runtime().block_on(async {
           let self_ref = self_ref.clone();
-          let mut memory_check_interval = time::interval(time::Duration::from_secs(2));
+          let memory_check_duration = if self_ref.server_config.server_params.read().worker_memory_limit.is_some(){
+            time::Duration::from_secs(15)
+          } else {
+            time::Duration::from_secs(60 * 60 * 24 * 365 * 100)
+          };
+
+          let mut memory_check_interval = time::interval(memory_check_duration);
+
+          self.invoke_hook("after_start");
+
           loop {
             tokio::select! {
               _ = receiver.changed() => {
@@ -308,7 +320,7 @@ impl ClusterMode {
                     if let Some(current_mem_usage) = largest_worker.memory_usage(){
                       if current_mem_usage > memory_limit {
                         largest_worker.reboot(self_ref.clone()).await.ok();
-                        if let Some(hook) = self_ref.server_config.server_params.read().hooks.get("after_memory_threshold_reached") {
+                        if let Some(hook) = self_ref.server_config.server_params.read().hooks.get("after_memory_limit_reached") {
                           call_with_gvl(|_|  hook.call::<_, Value>((largest_worker.pid(),)).ok() );
                         }
                       }

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

@@ -19,7 +19,7 @@ use itsi_rb_helpers::{
     call_with_gvl, call_without_gvl, create_ruby_thread, funcall_no_ret, print_rb_backtrace,
 };
 use itsi_tracing::{debug, error, info};
-use magnus::value::ReprValue;
+use magnus::{value::ReprValue, Value};
 use nix::unistd::Pid;
 use parking_lot::RwLock;
 use std::{
@@ -83,13 +83,13 @@ impl SingleMode {
         builder
             .thread_name("itsi-server-accept-loop")
             .thread_stack_size(3 * 1024 * 1024)
-            .enable_io()
-            .enable_time()
+            .enable_all()
             .build()
             .expect("Failed to build Tokio runtime")
     }
 
     pub fn stop(&self) -> Result<()> {
+        SHUTDOWN_REQUESTED.store(true, std::sync::atomic::Ordering::SeqCst);
         self.lifecycle_channel.send(LifecycleEvent::Shutdown).ok();
         Ok(())
     }
@@ -159,12 +159,12 @@ impl SingleMode {
     pub fn start_monitors(
         self: Arc<Self>,
         thread_workers: Arc<Vec<Arc<ThreadWorker>>>,
-    ) -> magnus::Thread {
+    ) -> Option<magnus::Thread> {
         call_with_gvl(move |_| {
             create_ruby_thread(move || {
                 call_without_gvl(move || {
                     let monitor_runtime = RuntimeBuilder::new_current_thread()
-                        .enable_time()
+                        .enable_all()
                         .build()
                         .unwrap();
                     let receiver = self.clone();
@@ -189,10 +189,10 @@ impl SingleMode {
                                   lifecycle_event = lifecycle_rx.recv() => {
                                       match lifecycle_event {
                                           Ok(LifecycleEvent::Restart) => {
-                                              receiver.restart().ok();
+                                              receiver.restart().await.ok();
                                           }
                                           Ok(LifecycleEvent::Reload) => {
-                                              receiver.reload().ok();
+                                              receiver.reload().await.ok();
                                           }
                                           Ok(LifecycleEvent::Shutdown) => {
                                             break;
@@ -231,14 +231,22 @@ impl SingleMode {
         );
 
         let (shutdown_sender, _) = watch::channel(RunningPhase::Running);
-        let thread = self.clone().start_monitors(thread_workers.clone());
+        let monitor_thread = self.clone().start_monitors(thread_workers.clone());
+        if monitor_thread.is_none() {
+            error!("Failed to start monitor thread");
+            return Err(ItsiError::new("Failed to start monitor thread"));
+        }
+        let monitor_thread = monitor_thread.unwrap();
         if SHUTDOWN_REQUESTED.load(Ordering::SeqCst) {
             return Ok(());
         }
-        runtime.block_on(
+        let result = runtime.block_on(
           async  {
               let server_params = self.server_config.server_params.read().clone();
-              server_params.middleware.get().unwrap().initialize_layers().await?;
+              if let Err(err) = server_params.initialize_middleware().await {
+                  error!("Failed to initialize middleware: {}", err);
+                  return Err(ItsiError::new("Failed to initialize middleware"))
+              }
               let tokio_listeners = server_params.listeners.lock()
                   .drain(..)
                   .map(|list| {
@@ -286,6 +294,7 @@ impl SingleMode {
                             }
                             lifecycle_event = lifecycle_rx.recv() => match lifecycle_event{
                               Ok(LifecycleEvent::Shutdown) => {
+                                debug!("Received lifecycle event: {:?}", lifecycle_event);
                                 shutdown_sender.send(RunningPhase::ShutdownPending).unwrap();
                                 tokio::time::sleep(Duration::from_millis(25)).await;
                                 for _i in 0..workers_clone.len() {
@@ -299,18 +308,42 @@ impl SingleMode {
                           }
                         }
                     }
-                    while let Some(_res) = acceptor_task_set.join_next().await {}
+
+                    let deadline = Instant::now()
+                        + Duration::from_secs_f64(self_ref.server_config.server_params.read().shutdown_timeout);
+                    tokio::select! {
+                        _ = async {
+                            while let Some(_res) = acceptor_task_set.join_next().await {}
+                        } => {},
+                        _ = tokio::time::sleep_until(tokio::time::Instant::from_std(deadline)) => {},
+                    }
                 });
 
               }
 
+              if self.is_single_mode() {
+                self.invoke_hook("after_start");
+              }
+
               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>(())
-          })?;
+          });
+
+        debug!("Single mode runtime exited.");
+
+        if result.is_err() {
+            for _i in 0..thread_workers.len() {
+                job_sender.send_blocking(RequestJob::Shutdown).unwrap();
+                nonblocking_sender
+                    .send_blocking(RequestJob::Shutdown)
+                    .unwrap();
+            }
+            self.lifecycle_channel.send(LifecycleEvent::Shutdown).ok();
+        }
 
         shutdown_sender.send(RunningPhase::Shutdown).ok();
         let deadline = Instant::now()
@@ -318,24 +351,34 @@ impl SingleMode {
 
         runtime.shutdown_timeout(Duration::from_millis(100));
 
+        debug!("Shutdown timeout finished.");
         loop {
             if thread_workers
                 .iter()
                 .all(|worker| call_with_gvl(move |_| !worker.poll_shutdown(deadline)))
             {
-                funcall_no_ret(thread, "join", ()).ok();
+                funcall_no_ret(monitor_thread, "join", ()).ok();
                 break;
             }
             sleep(Duration::from_millis(50));
         }
 
+        if self.is_single_mode() {
+            self.invoke_hook("before_shutdown");
+        }
+
         if self.restart_requested.load(Ordering::SeqCst) {
             self.restart_requested.store(false, Ordering::SeqCst);
             info!("Worker restarting");
             self.run()?;
         }
+
         debug!("Runtime has shut down");
-        Ok(())
+        result
+    }
+
+    pub fn is_single_mode(&self) -> bool {
+        self.server_config.server_params.read().workers == 1
     }
 
     pub(crate) async fn serve_connection(
@@ -367,7 +410,7 @@ impl SingleMode {
         let mut serve = Box::pin(
             binding
                 .timer(TokioTimer::new())
-                .header_read_timeout(Duration::from_secs(1))
+                .header_read_timeout(self.server_config.server_params.read().header_read_timeout)
                 .serve_connection_with_upgrades(io, service),
         );
 
@@ -400,9 +443,15 @@ impl SingleMode {
     /// Attempts to reload the config "live"
     /// Not that when running in single mode this will not unload
     /// old code. If you need a clean restart, use the `restart` (SIGHUP) method instead
-    pub fn reload(&self) -> Result<()> {
+    pub async fn reload(&self) -> Result<()> {
+        if !self.server_config.check_config().await {
+            return Ok(());
+        }
         let should_reexec = self.server_config.clone().reload(false)?;
         if should_reexec {
+            if self.is_single_mode() {
+                self.invoke_hook("before_restart");
+            }
             self.server_config.dup_fds()?;
             self.server_config.reload_exec()?;
         }
@@ -412,8 +461,19 @@ impl SingleMode {
         Ok(())
     }
 
+    pub fn invoke_hook(&self, hook_name: &str) {
+        if let Some(hook) = self.server_config.server_params.read().hooks.get(hook_name) {
+            call_with_gvl(|_| hook.call::<_, Value>(()).ok());
+        }
+    }
     /// Restart the server while keeping connections open.
-    pub fn restart(&self) -> Result<()> {
+    pub async fn restart(&self) -> Result<()> {
+        if !self.server_config.check_config().await {
+            return Ok(());
+        }
+        if self.is_single_mode() {
+            self.invoke_hook("before_restart");
+        }
         self.server_config.dup_fds()?;
         self.server_config.reload_exec()?;
         Ok(())

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

@@ -49,6 +49,7 @@ fn receive_signal(signum: i32, _: sighandler_t) {
 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

@@ -10,6 +10,7 @@ use std::sync::atomic::AtomicUsize;
 use std::sync::atomic::Ordering;
 use std::task::Context;
 use std::task::Poll;
+use tracing::debug;
 
 /// Custom error to indicate that the maximum body size was exceeded.
 #[derive(Debug)]
@@ -71,6 +72,11 @@ where
                         Ok(data) => {
                             let len = data.remaining();
                             self.current += len;
+                            debug!(
+                              target: "option::max_body",
+                              "current: {}, limit: {}",
+                              self.current, self.limit.load(Ordering::Relaxed)
+                            );
                             if self.current > self.limit.load(Ordering::Relaxed) {
                                 Poll::Ready(Some(Err(Box::new(MaxBodySizeReached))))
                             } else {

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

@@ -1,4 +1,5 @@
 use async_channel::Sender;
+use itsi_error::ItsiError;
 use itsi_rb_helpers::{
     call_with_gvl, call_without_gvl, create_ruby_thread, kill_threads, HeapValue,
 };
@@ -159,11 +160,10 @@ impl ThreadWorker {
     pub fn poll_shutdown(&self, deadline: Instant) -> bool {
         if let Some(thread) = self.thread.read().deref() {
             if Instant::now() > deadline {
-                warn!("Worker shutdown timed out. Killing thread");
+                debug!("Worker shutdown timed out. Killing thread {:?}", thread);
                 self.terminated.store(true, Ordering::SeqCst);
                 kill_threads(vec![thread.as_value()]);
             }
-            debug!("Checking thread status");
             if thread.funcall::<_, _, bool>(*ID_ALIVE, ()).unwrap_or(false) {
                 return true;
             }
@@ -184,6 +184,7 @@ impl ThreadWorker {
         call_with_gvl(|_| {
             *self.thread.write() = Some(
                 create_ruby_thread(move || {
+                    debug!("Ruby thread worker started");
                     if let Some(scheduler_class) = scheduler_class {
                         if let Err(err) = self_ref.fiber_accept_loop(
                             params,
@@ -198,6 +199,9 @@ impl ThreadWorker {
                         self_ref.accept_loop(params, name, receiver, terminated);
                     }
                 })
+                .ok_or_else(|| {
+                    ItsiError::InternalServerError("Failed to create Ruby thread".to_owned())
+                })?
                 .into(),
             );
             Ok::<(), magnus::Error>(())

data/crates/itsi_server/src/services/itsi_http_service.rs

@@ -1,5 +1,5 @@
 use crate::default_responses::{NOT_FOUND_RESPONSE, TIMEOUT_RESPONSE};
-use crate::ruby_types::itsi_server::itsi_server_config::ServerParams;
+use crate::ruby_types::itsi_server::itsi_server_config::{ItsiServerTokenPreference, ServerParams};
 use crate::server::binds::listener::ListenerInfo;
 use crate::server::http_message_types::{ConversionExt, HttpResponse, RequestExt, ResponseFormat};
 use crate::server::lifecycle_event::LifecycleEvent;
@@ -18,6 +18,7 @@ use itsi_error::ItsiError;
 use regex::Regex;
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::OnceLock;
+use tracing::error;
 
 use std::{future::Future, ops::Deref, pin::Pin, sync::Arc};
 use tokio::sync::watch::{self};
@@ -153,6 +154,10 @@ impl HttpRequestContext {
     }
 }
 
+const SERVER_TOKEN_VERSION: HeaderValue =
+    HeaderValue::from_static(concat!("Itsi/", env!("CARGO_PKG_VERSION")));
+const SERVER_TOKEN_NAME: HeaderValue = HeaderValue::from_static("Itsi");
+
 impl Service<Request<Incoming>> for ItsiHttpService {
     type Response = HttpResponse;
     type Error = ItsiError;
@@ -195,7 +200,10 @@ impl Service<Request<Incoming>> for ItsiHttpService {
                         depth = index;
                         break;
                     }
-                    Err(e) => return Err(e.into()),
+                    Err(e) => {
+                        error!("Middleware error: {}", e);
+                        break;
+                    }
                 }
             }
 
@@ -208,6 +216,16 @@ impl Service<Request<Incoming>> for ItsiHttpService {
                 resp = elm.after(resp, &mut context).await;
             }
 
+            match params.itsi_server_token_preference {
+                ItsiServerTokenPreference::Version => {
+                    resp.headers_mut().insert("Server", SERVER_TOKEN_VERSION);
+                }
+                ItsiServerTokenPreference::Name => {
+                    resp.headers_mut().insert("Server", SERVER_TOKEN_NAME);
+                }
+                ItsiServerTokenPreference::None => {}
+            }
+
             Ok(resp)
         };
 

data/crates/itsi_server/src/services/rate_limiter.rs

@@ -5,10 +5,12 @@ use redis::{Client, RedisError, Script};
 use serde::Deserialize;
 use std::any::Any;
 use std::collections::{HashMap, HashSet};
+use std::result::Result;
 use std::sync::{Arc, LazyLock, Mutex};
 use std::time::{Duration, Instant};
 use tokio::sync::{Mutex as AsyncMutex, RwLock};
 use tokio::time::timeout;
+use tracing::warn;
 use url::Url;
 
 #[derive(Debug)]
@@ -109,10 +111,19 @@ impl RedisRateLimiter {
         // Create the Lua script once when initializing the rate limiter
         let increment_script = Script::new(
             r#"
+            -- Increment the counter
             local current = redis.call('INCR', KEYS[1])
-            if redis.call('TTL', KEYS[1]) < 0 then
-                redis.call('EXPIRE', KEYS[1], ARGV[1])
+
+            -- Fetch existing TTL
+            local ttl = redis.call('TTL', KEYS[1])
+            if ttl < 0 then
+                -- If no TTL is set, apply the window interval
+                local window = tonumber(ARGV[1])
+                redis.call('EXPIRE', KEYS[1], window)
+                ttl = window
             end
+
+            -- Return both the current count and remaining TTL
             return { current, ttl }
             "#,
         );
@@ -135,14 +146,14 @@ impl RedisRateLimiter {
         let mut connection = (*self.connection).clone();
 
         // Set the ban with the reason as the value
-        let _: () = redis::cmd("SET")
+        let _: Result<(), RedisError> = redis::cmd("SET")
             .arg(&ban_key)
             .arg(reason)
             .arg("EX")
             .arg(timeout_secs)
             .query_async(&mut connection)
             .await
-            .map_err(RateLimitError::RedisError)?;
+            .inspect_err(|e| warn!("Exception banning IP {:?}", e));
 
         Ok(())
     }