itsi-scheduler 0.2.15 → 0.2.16

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

@@ -3,6 +3,7 @@ use either::Either;
 use itsi_tracing::*;
 use magnus::error::Result;
 use serde::Deserialize;
+use tracing::enabled;
 
 use crate::server::http_message_types::{HttpRequest, HttpResponse};
 use crate::services::itsi_http_service::HttpRequestContext;
@@ -39,18 +40,6 @@ pub enum LogMiddlewareLevel {
     Error,
 }
 
-impl LogMiddlewareLevel {
-    pub fn log(&self, message: String) {
-        match self {
-            LogMiddlewareLevel::Trace => trace!(target: "middleware::log_requests", message),
-            LogMiddlewareLevel::Debug => debug!(target: "middleware::log_requests", message),
-            LogMiddlewareLevel::Info => info!(target: "middleware::log_requests", message),
-            LogMiddlewareLevel::Warn => warn!(target: "middleware::log_requests", message),
-            LogMiddlewareLevel::Error => error!(target: "middleware::log_requests", message),
-        }
-    }
-}
-
 #[async_trait]
 impl MiddlewareLayer for LogRequests {
     async fn initialize(&self) -> Result<()> {
@@ -64,7 +53,38 @@ impl MiddlewareLayer for LogRequests {
     ) -> Result<Either<HttpRequest, HttpResponse>> {
         context.init_logging_params();
         if let Some(LogConfig { level, format }) = self.before.as_ref() {
-            level.log(format.rewrite_request(&req, context));
+            match level {
+                LogMiddlewareLevel::Trace => {
+                    if enabled!(target: "middleware::log_requests", tracing::Level::TRACE) {
+                        let message = format.rewrite_request(&req, context);
+                        trace!(target: "middleware::log_requests", message);
+                    }
+                }
+                LogMiddlewareLevel::Debug => {
+                    if enabled!(target: "middleware::log_requests", tracing::Level::DEBUG) {
+                        let message = format.rewrite_request(&req, context);
+                        debug!(target: "middleware::log_requests", message);
+                    }
+                }
+                LogMiddlewareLevel::Info => {
+                    if enabled!(target: "middleware::log_requests", tracing::Level::INFO) {
+                        let message = format.rewrite_request(&req, context);
+                        info!(target: "middleware::log_requests", message);
+                    }
+                }
+                LogMiddlewareLevel::Warn => {
+                    if enabled!(target: "middleware::log_requests", tracing::Level::WARN) {
+                        let message = format.rewrite_request(&req, context);
+                        warn!(target: "middleware::log_requests", message);
+                    }
+                }
+                LogMiddlewareLevel::Error => {
+                    if enabled!(target: "middleware::log_requests", tracing::Level::ERROR) {
+                        let message = format.rewrite_request(&req, context);
+                        error!(target: "middleware::log_requests", message);
+                    }
+                }
+            }
         }
 
         Ok(Either::Left(req))
@@ -72,7 +92,38 @@ impl MiddlewareLayer for LogRequests {
 
     async fn after(&self, resp: HttpResponse, context: &mut HttpRequestContext) -> HttpResponse {
         if let Some(LogConfig { level, format }) = self.after.as_ref() {
-            level.log(format.rewrite_response(&resp, context));
+            match level {
+                LogMiddlewareLevel::Trace => {
+                    if enabled!(target: "middleware::log_requests", tracing::Level::TRACE) {
+                        let message = format.rewrite_response(&resp, context);
+                        trace!(target: "middleware::log_requests", message);
+                    }
+                }
+                LogMiddlewareLevel::Debug => {
+                    if enabled!(target: "middleware::log_requests", tracing::Level::DEBUG) {
+                        let message = format.rewrite_response(&resp, context);
+                        debug!(target: "middleware::log_requests", message);
+                    }
+                }
+                LogMiddlewareLevel::Info => {
+                    if enabled!(target: "middleware::log_requests", tracing::Level::INFO) {
+                        let message = format.rewrite_response(&resp, context);
+                        info!(target: "middleware::log_requests", message);
+                    }
+                }
+                LogMiddlewareLevel::Warn => {
+                    if enabled!(target: "middleware::log_requests", tracing::Level::WARN) {
+                        let message = format.rewrite_response(&resp, context);
+                        warn!(target: "middleware::log_requests", message);
+                    }
+                }
+                LogMiddlewareLevel::Error => {
+                    if enabled!(target: "middleware::log_requests", tracing::Level::ERROR) {
+                        let message = format.rewrite_response(&resp, context);
+                        error!(target: "middleware::log_requests", message);
+                    }
+                }
+            }
         }
 
         resp

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

@@ -37,7 +37,7 @@ impl MiddlewareLayer for MaxBody {
     async fn after(&self, resp: HttpResponse, context: &mut HttpRequestContext) -> HttpResponse {
         if resp.status() == StatusCode::PAYLOAD_TOO_LARGE {
             self.error_response
-                .to_http_response(context.response_format().clone())
+                .to_http_response(*context.response_format())
                 .await
         } else {
             resp

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

@@ -303,7 +303,7 @@ impl MiddlewareLayer for Proxy {
         let url = self.to.rewrite_request(&req, context);
 
         let accept: ResponseFormat = req.accept().into();
-        let error_response = self.error_response.to_http_response(accept.clone()).await;
+        let error_response = self.error_response.to_http_response(accept).await;
 
         let destination = match Url::parse(&url) {
             Ok(dest) => dest,

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

@@ -6,6 +6,7 @@ use crate::services::rate_limiter::{
 };
 use async_trait::async_trait;
 use either::Either;
+use http::{HeaderName, HeaderValue};
 use magnus::error::Result;
 use serde::Deserialize;
 use std::collections::HashMap;
@@ -24,6 +25,8 @@ pub struct RateLimit {
     pub trusted_proxies: HashMap<String, TokenSource>,
     #[serde(default = "too_many_requests_error_response")]
     pub error_response: ErrorResponse,
+    #[serde(skip)]
+    pub limit_header_value: OnceLock<HeaderValue>,
 }
 
 fn too_many_requests_error_response() -> ErrorResponse {
@@ -38,6 +41,12 @@ pub enum RateLimitKey {
     Parameter(TokenSource),
 }
 
+static X_RATELIMIT_LIMIT: HeaderName = HeaderName::from_static("x-ratelimit-limit");
+static X_RATELIMIT_REMAINING: HeaderName = HeaderName::from_static("x-ratelimit-remaining");
+static X_RATELIMIT_RESET: HeaderName = HeaderName::from_static("x-ratelimit-reset");
+static RETRY_AFTER: HeaderName = HeaderName::from_static("retry-after");
+static ZERO_VALUE: HeaderValue = HeaderValue::from_static("0");
+
 #[async_trait]
 impl MiddlewareLayer for RateLimit {
     async fn initialize(&self) -> Result<()> {
@@ -46,6 +55,9 @@ impl MiddlewareLayer for RateLimit {
         if let Ok(limiter) = get_rate_limiter(&self.store_config).await {
             let _ = self.rate_limiter.set(limiter);
         }
+        self.limit_header_value
+            .set(self.requests.to_string().parse().unwrap())
+            .ok();
         Ok(())
     }
 
@@ -58,8 +70,7 @@ impl MiddlewareLayer for RateLimit {
         let key_value = match &self.key {
             RateLimitKey::SocketAddress => {
                 // Use the socket address from the context
-                if self.trusted_proxies.contains_key(&context.addr) {
-                    let source = self.trusted_proxies.get(&context.addr).unwrap();
+                if let Some(source) = self.trusted_proxies.get(&context.addr) {
                     source.extract_token(&req).unwrap_or(&context.addr)
                 } else {
                     &context.addr
@@ -114,18 +125,20 @@ impl MiddlewareLayer for RateLimit {
                         .error_response
                         .to_http_response(req.accept().into())
                         .await;
+                    let ttl_header_value: HeaderValue = ttl.to_string().parse().unwrap();
+                    response.headers_mut().insert(
+                        X_RATELIMIT_LIMIT.clone(),
+                        self.limit_header_value.get().unwrap().clone(),
+                    );
                     response
                         .headers_mut()
-                        .insert("X-RateLimit-Limit", limit.to_string().parse().unwrap());
-                    response
-                        .headers_mut()
-                        .insert("X-RateLimit-Remaining", "0".parse().unwrap());
+                        .insert(X_RATELIMIT_REMAINING.clone(), ZERO_VALUE.clone());
                     response
                         .headers_mut()
-                        .insert("X-RateLimit-Reset", ttl.to_string().parse().unwrap());
+                        .insert(X_RATELIMIT_RESET.clone(), ttl_header_value.clone());
                     response
                         .headers_mut()
-                        .insert("Retry-After", ttl.to_string().parse().unwrap());
+                        .insert(RETRY_AFTER.clone(), ttl_header_value);
                     Ok(Either::Right(response))
                 }
                 Err(e) => {

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

@@ -119,11 +119,7 @@ impl MiddlewareLayer for RubyApp {
         if self.sendfile {
             if let Some(sendfile_header) = resp.headers().get("X-Sendfile") {
                 return ROOT_STATIC_FILE_SERVER
-                    .serve_single_abs(
-                        sendfile_header.to_str().unwrap(),
-                        context.accept.clone(),
-                        &[],
-                    )
+                    .serve_single_abs(sendfile_header.to_str().unwrap(), context.accept, &[])
                     .await;
             }
         }

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

@@ -21,6 +21,10 @@ pub struct StaticResponse {
     body: Vec<u8>,
     #[serde(skip)]
     header_map: OnceLock<HeaderMap>,
+    #[serde(skip)]
+    body_bytes: OnceLock<Full<Bytes>>,
+    #[serde(skip)]
+    status_code: OnceLock<StatusCode>,
 }
 
 #[async_trait]
@@ -35,6 +39,12 @@ impl MiddlewareLayer for StaticResponse {
         self.header_map
             .set(header_map)
             .map_err(|_| ItsiError::new("Failed to set headers"))?;
+        self.body_bytes
+            .set(Full::new(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 +53,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(BoxBody::new(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/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))]

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

@@ -0,0 +1,96 @@
+use std::{ops::Deref, pin::Pin, sync::Arc, time::Duration};
+
+use hyper_util::rt::TokioIo;
+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();
+        self.join_set.spawn(async move {
+            let executor = &acceptor_args.strategy.executor;
+            let mut serve = Box::pin(executor.serve_connection_with_upgrades(
+                io,
+                ItsiHttpService {
+                    inner: Arc::new(ItsiHttpServiceInner {
+                        acceptor_args: acceptor_args.clone(),
+                        addr: addr.to_string(),
+                    }),
+                },
+            ));
+
+            tokio::select! {
+                // Await the connection finishing naturally.
+                res = &mut serve => {
+                    match res {
+                        Ok(()) => {
+                            debug!("Connection closed normally");
+                        },
+                        Err(res) => {
+                            debug!("Connection closed abruptly: {:?}", res);
+                        }
+                    }
+                    serve.as_mut().graceful_shutdown();
+                },
+                // 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 => {
+                debug!("Shutdown timeout reached; abandoning remaining acceptor tasks.");
+            }
+        }
+    }
+}

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;