itsi-server 0.2.15 → 0.2.17

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

@@ -3,13 +3,12 @@ use itsi_error::ItsiError;
 use itsi_rb_helpers::{
     call_with_gvl, call_without_gvl, create_ruby_thread, kill_threads, HeapValue,
 };
-use itsi_tracing::{debug, error, warn};
+use itsi_tracing::{debug, error};
 use magnus::{
     error::Result,
     value::{InnerValue, Lazy, LazyId, Opaque, ReprValue},
     Module, RClass, Ruby, Thread, Value,
 };
-use nix::unistd::Pid;
 use parking_lot::{Mutex, RwLock};
 use std::{
     ops::Deref,
@@ -17,22 +16,24 @@ use std::{
         atomic::{AtomicBool, AtomicU64, Ordering},
         Arc,
     },
-    thread,
-    time::{Duration, Instant, SystemTime, UNIX_EPOCH},
+    time::{Instant, SystemTime, UNIX_EPOCH},
 };
 use tokio::{runtime::Builder as RuntimeBuilder, sync::watch};
 use tracing::instrument;
 
-use crate::ruby_types::{
-    itsi_grpc_call::ItsiGrpcCall, itsi_http_request::ItsiHttpRequest,
-    itsi_server::itsi_server_config::ServerParams, ITSI_SERVER,
+use crate::{
+    ruby_types::{
+        itsi_grpc_call::ItsiGrpcCall, itsi_http_request::ItsiHttpRequest,
+        itsi_server::itsi_server_config::ServerParams, ITSI_SERVER,
+    },
+    server::process_worker::CORE_IDS,
 };
 
 use super::request_job::RequestJob;
 pub struct ThreadWorker {
     pub params: Arc<ServerParams>,
     pub id: u8,
-    pub name: String,
+    pub worker_id: usize,
     pub request_id: AtomicU64,
     pub current_request_start: AtomicU64,
     pub receiver: Arc<async_channel::Receiver<RequestJob>>,
@@ -61,8 +62,11 @@ type ThreadWorkerBuildResult = Result<(
     Sender<RequestJob>,
 )>;
 
-#[instrument(name = "boot", parent=None, skip(params, pid))]
-pub fn build_thread_workers(params: Arc<ServerParams>, pid: Pid) -> ThreadWorkerBuildResult {
+#[instrument(name = "boot", parent=None, skip(params, worker_id))]
+pub fn build_thread_workers(
+    params: Arc<ServerParams>,
+    worker_id: usize,
+) -> ThreadWorkerBuildResult {
     let blocking_thread_count = params.threads;
     let nonblocking_thread_count = params.scheduler_threads;
     let ruby_thread_request_backlog_size: usize = params
@@ -80,7 +84,7 @@ pub fn build_thread_workers(params: Arc<ServerParams>, pid: Pid) -> ThreadWorker
             ThreadWorker::new(
                 params.clone(),
                 id,
-                format!("{:?}#{:?}", pid, id),
+                worker_id,
                 blocking_receiver_ref.clone(),
                 blocking_sender_ref.clone(),
                 if nonblocking_thread_count.is_some() {
@@ -103,7 +107,7 @@ pub fn build_thread_workers(params: Arc<ServerParams>, pid: Pid) -> ThreadWorker
             workers.push(ThreadWorker::new(
                 params.clone(),
                 id,
-                format!("{:?}#{:?}", pid, id),
+                worker_id,
                 nonblocking_receiver_ref.clone(),
                 nonblocking_sender_ref.clone(),
                 Some(scheduler_class),
@@ -138,7 +142,7 @@ impl ThreadWorker {
     pub fn new(
         params: Arc<ServerParams>,
         id: u8,
-        name: String,
+        worker_id: usize,
         receiver: Arc<async_channel::Receiver<RequestJob>>,
         sender: Sender<RequestJob>,
         scheduler_class: Option<Opaque<Value>>,
@@ -146,9 +150,9 @@ impl ThreadWorker {
         let worker = Arc::new(Self {
             params,
             id,
+            worker_id,
             request_id: AtomicU64::new(0),
             current_request_start: AtomicU64::new(0),
-            name,
             receiver,
             sender,
             thread: RwLock::new(None),
@@ -178,20 +182,24 @@ impl ThreadWorker {
     }
 
     pub fn run(self: Arc<Self>) -> Result<()> {
-        let name = self.name.clone();
         let receiver = self.receiver.clone();
         let terminated = self.terminated.clone();
         let scheduler_class = self.scheduler_class;
         let params = self.params.clone();
         let self_ref = self.clone();
+        let worker_id = self.worker_id;
         call_with_gvl(|_| {
             *self.thread.write() = Some(
                 create_ruby_thread(move || {
+                    if params.pin_worker_cores {
+                        core_affinity::set_for_current(
+                            CORE_IDS[((2 * worker_id) + 1) % CORE_IDS.len()],
+                        );
+                    }
                     debug!("Ruby thread worker started");
                     if let Some(scheduler_class) = scheduler_class {
                         if let Err(err) = self_ref.fiber_accept_loop(
                             params,
-                            name,
                             receiver,
                             scheduler_class,
                             terminated,
@@ -199,7 +207,7 @@ impl ThreadWorker {
                             error!("Error in fiber_accept_loop: {:?}", err);
                         }
                     } else {
-                        self_ref.accept_loop(params, name, receiver, terminated);
+                        self_ref.accept_loop(params, receiver, terminated);
                     }
                 })
                 .ok_or_else(|| {
@@ -255,9 +263,14 @@ impl ThreadWorker {
                         }
                     }
                 }
+
                 for _ in 0..MAX_BATCH_SIZE {
                     if let Ok(req) = receiver.try_recv() {
+                        let should_break = matches!(req, RequestJob::Shutdown);
                         batch.push(req);
+                        if should_break {
+                            break;
+                        }
                     } else {
                         break;
                     }
@@ -300,7 +313,9 @@ impl ThreadWorker {
                                     ItsiGrpcCall::internal_error(ruby, response, err)
                                 }
                             }
-                            RequestJob::Shutdown => return true,
+                            RequestJob::Shutdown => {
+                                return true;
+                            }
                         }
                     }
                     false
@@ -332,15 +347,14 @@ impl ThreadWorker {
                 if yield_result.is_err() {
                     break;
                 }
-            })
+            });
         })
     }
 
-    #[instrument(skip_all, fields(thread_worker=name))]
+    #[instrument(skip_all, fields(thread_worker=format!("{}:{}", self.id, self.worker_id)))]
     pub fn fiber_accept_loop(
         self: Arc<Self>,
         params: Arc<ServerParams>,
-        name: String,
         receiver: Arc<async_channel::Receiver<RequestJob>>,
         scheduler_class: Opaque<Value>,
         terminated: Arc<AtomicBool>,
@@ -415,68 +429,76 @@ impl ThreadWorker {
         });
     }
 
-    #[instrument(skip_all, fields(thread_worker=id))]
+    #[instrument(skip_all, fields(thread_worker=format!("{}:{}", self.id, self.worker_id)))]
     pub fn accept_loop(
         self: Arc<Self>,
         params: Arc<ServerParams>,
-        id: String,
         receiver: Arc<async_channel::Receiver<RequestJob>>,
         terminated: Arc<AtomicBool>,
     ) {
-        let ruby = Ruby::get().unwrap();
         let mut idle_counter = 0;
-        let self_ref = self.clone();
         call_without_gvl(|| loop {
-            if receiver.is_empty() {
-                if let Some(oob_gc_threshold) = params.oob_gc_responses_threshold {
-                    idle_counter = (idle_counter + 1) % oob_gc_threshold;
-                    if idle_counter == 0 {
-                        call_with_gvl(|_ruby| {
-                            ruby.gc_start();
-                        });
-                    }
-                };
-            }
             match receiver.recv_blocking() {
-                Ok(RequestJob::ProcessHttpRequest(request, app_proc)) => {
-                    self_ref.request_id.fetch_add(1, Ordering::Relaxed);
-                    self_ref.current_request_start.store(
-                        SystemTime::now()
-                            .duration_since(UNIX_EPOCH)
-                            .unwrap()
-                            .as_secs(),
-                        Ordering::Relaxed,
-                    );
-                    call_with_gvl(|_ruby| {
-                        request.process(&ruby, app_proc).ok();
-                    });
-                    if terminated.load(Ordering::Relaxed) {
-                        break;
+                Err(_) => break,
+                Ok(RequestJob::Shutdown) => break,
+                Ok(request_job) => call_with_gvl(|ruby| {
+                    self.process_one(&ruby, request_job, &terminated);
+                    while let Ok(request_job) = receiver.try_recv() {
+                        if matches!(request_job, RequestJob::Shutdown) {
+                            terminated.store(true, Ordering::Relaxed);
+                            break;
+                        }
+                        self.process_one(&ruby, request_job, &terminated);
                     }
-                }
-                Ok(RequestJob::ProcessGrpcRequest(request, app_proc)) => {
-                    self_ref.request_id.fetch_add(1, Ordering::Relaxed);
-                    self_ref.current_request_start.store(
-                        SystemTime::now()
-                            .duration_since(UNIX_EPOCH)
-                            .unwrap()
-                            .as_secs(),
-                        Ordering::Relaxed,
-                    );
-                    call_with_gvl(|_ruby| {
-                        request.process(&ruby, app_proc).ok();
-                    });
-                    if terminated.load(Ordering::Relaxed) {
-                        break;
+                    if let Some(thresh) = params.oob_gc_responses_threshold {
+                        idle_counter = (idle_counter + 1) % thresh;
+                        if idle_counter == 0 {
+                            ruby.gc_start();
+                        }
                     }
+                }),
+            };
+            if terminated.load(Ordering::Relaxed) {
+                break;
+            }
+        });
+    }
+
+    fn process_one(self: &Arc<Self>, ruby: &Ruby, job: RequestJob, terminated: &Arc<AtomicBool>) {
+        match job {
+            RequestJob::ProcessHttpRequest(request, app_proc) => {
+                if terminated.load(Ordering::Relaxed) {
+                    request.response().unwrap().service_unavailable();
+                    return;
                 }
-                Ok(RequestJob::Shutdown) => {
-                    break;
-                }
-                Err(_) => {
-                    thread::sleep(Duration::from_micros(1));
+                self.request_id.fetch_add(1, Ordering::Relaxed);
+                self.current_request_start.store(
+                    SystemTime::now()
+                        .duration_since(UNIX_EPOCH)
+                        .unwrap()
+                        .as_secs(),
+                    Ordering::Relaxed,
+                );
+                request.process(ruby, app_proc).ok();
+            }
+
+            RequestJob::ProcessGrpcRequest(request, app_proc) => {
+                if terminated.load(Ordering::Relaxed) {
+                    request.stream().unwrap().close().ok();
+                    return;
                 }
+                self.request_id.fetch_add(1, Ordering::Relaxed);
+                self.current_request_start.store(
+                    SystemTime::now()
+                        .duration_since(UNIX_EPOCH)
+                        .unwrap()
+                        .as_secs(),
+                    Ordering::Relaxed,
+                );
+                request.process(ruby, app_proc).ok();
             }
-        });
+
+            RequestJob::Shutdown => unreachable!(),
+        }
     }
 }

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

@@ -1,30 +1,25 @@
 use crate::default_responses::{NOT_FOUND_RESPONSE, TIMEOUT_RESPONSE};
-use crate::ruby_types::itsi_server::itsi_server_config::{ItsiServerTokenPreference, ServerParams};
-use crate::server::binds::listener::ListenerInfo;
+use crate::ruby_types::itsi_server::itsi_server_config::ItsiServerTokenPreference;
 use crate::server::http_message_types::{
     ConversionExt, HttpRequest, HttpResponse, RequestExt, ResponseFormat,
 };
 use crate::server::lifecycle_event::LifecycleEvent;
 use crate::server::middleware_stack::MiddlewareLayer;
-use crate::server::request_job::RequestJob;
-use crate::server::serve_strategy::single_mode::RunningPhase;
-use crate::server::signal::send_lifecycle_event;
+use crate::server::serve_strategy::acceptor::AcceptorArgs;
+use crate::server::signal::{send_lifecycle_event, SHUTDOWN_REQUESTED};
 use chrono::{self, DateTime, Local};
 use either::Either;
 use http::header::ACCEPT_ENCODING;
 use http::{HeaderValue, Request};
 use hyper::body::Incoming;
-use hyper::service::Service;
-use itsi_error::ItsiError;
 use regex::Regex;
+use smallvec::SmallVec;
+use std::ops::Deref;
 use std::sync::atomic::{AtomicBool, Ordering};
-use std::sync::OnceLock;
+use std::sync::{Arc, OnceLock};
 use std::time::{Duration, Instant};
-use tracing::error;
-
-use std::{future::Future, ops::Deref, pin::Pin, sync::Arc};
-use tokio::sync::watch::{self};
 use tokio::time::timeout;
+use tracing::error;
 
 #[derive(Clone)]
 pub struct ItsiHttpService {
@@ -40,12 +35,16 @@ impl Deref for ItsiHttpService {
 }
 
 pub struct ItsiHttpServiceInner {
-    pub sender: async_channel::Sender<RequestJob>,
-    pub nonblocking_sender: async_channel::Sender<RequestJob>,
-    pub server_params: Arc<ServerParams>,
-    pub listener: Arc<ListenerInfo>,
+    pub acceptor_args: Arc<AcceptorArgs>,
     pub addr: String,
-    pub shutdown_channel: watch::Receiver<RunningPhase>,
+}
+
+impl Deref for ItsiHttpServiceInner {
+    type Target = Arc<AcceptorArgs>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.acceptor_args
+    }
 }
 
 #[derive(Clone)]
@@ -79,12 +78,14 @@ pub struct RequestContextInner {
     pub request_start_time: OnceLock<DateTime<Local>>,
     pub start_instant: Instant,
     pub if_none_match: OnceLock<Option<String>>,
-    pub supported_encoding_set: OnceLock<Vec<HeaderValue>>,
+    pub supported_encoding_set: OnceLock<AcceptEncodingSet>,
     pub is_ruby_request: Arc<AtomicBool>,
 }
 
+type AcceptEncodingSet = SmallVec<[HeaderValue; 2]>;
+
 impl HttpRequestContext {
-    fn new(
+    pub fn new(
         service: ItsiHttpService,
         matching_pattern: Option<Arc<Regex>>,
         accept: ResponseFormat,
@@ -108,12 +109,14 @@ impl HttpRequestContext {
     }
 
     pub fn set_supported_encoding_set(&self, req: &HttpRequest) {
-        self.inner.supported_encoding_set.get_or_init(move || {
-            req.headers()
-                .get_all(ACCEPT_ENCODING)
-                .into_iter()
-                .cloned()
-                .collect::<Vec<_>>()
+        self.inner.supported_encoding_set.get_or_init(|| {
+            let mut set: AcceptEncodingSet = SmallVec::new();
+
+            for hv in req.headers().get_all(ACCEPT_ENCODING) {
+                set.push(hv.clone()); // clone ≈ 16 B struct copy
+            }
+
+            set
         });
     }
 
@@ -163,7 +166,7 @@ impl HttpRequestContext {
         self.inner.response_format.get().unwrap()
     }
 
-    pub fn supported_encoding_set(&self) -> Option<&Vec<HeaderValue>> {
+    pub fn supported_encoding_set(&self) -> Option<&AcceptEncodingSet> {
         self.inner.supported_encoding_set.get()
     }
 }
@@ -172,32 +175,31 @@ 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;
-    type Future = Pin<Box<dyn Future<Output = itsi_error::Result<HttpResponse>> + Send>>;
-
-    fn call(&self, req: Request<Incoming>) -> Self::Future {
-        let params = self.server_params.clone();
-        let self_clone = self.clone();
+impl ItsiHttpService {
+    pub async fn handle_request(&self, req: Request<Incoming>) -> itsi_error::Result<HttpResponse> {
         let mut req = req.limit();
         let accept: ResponseFormat = req.accept().into();
-        let accept_clone = accept.clone();
         let is_single_mode = self.server_params.workers == 1;
 
         let request_timeout = self.server_params.request_timeout;
         let is_ruby_request = Arc::new(AtomicBool::new(false));
         let irr_clone = is_ruby_request.clone();
+
+        let token_preference = self.server_params.itsi_server_token_preference;
+
         let service_future = async move {
+            let middleware_stack = self
+                .server_params
+                .middleware
+                .get()
+                .unwrap()
+                .stack_for(&req)
+                .unwrap();
+            let (stack, matching_pattern) = middleware_stack;
             let mut resp: Option<HttpResponse> = None;
-            let (stack, matching_pattern) = params.middleware.get().unwrap().stack_for(&req)?;
 
-            let mut context = HttpRequestContext::new(
-                self_clone,
-                matching_pattern,
-                accept_clone.clone(),
-                irr_clone,
-            );
+            let mut context =
+                HttpRequestContext::new(self.clone(), matching_pattern, accept, irr_clone);
             let mut depth = 0;
 
             for (index, elm) in stack.iter().enumerate() {
@@ -217,14 +219,14 @@ impl Service<Request<Incoming>> for ItsiHttpService {
 
             let mut resp = match resp {
                 Some(r) => r,
-                None => return Ok(NOT_FOUND_RESPONSE.to_http_response(accept_clone).await),
+                None => return Ok(NOT_FOUND_RESPONSE.to_http_response(accept).await),
             };
 
             for elm in stack.iter().rev().skip(stack.len() - depth - 1) {
                 resp = elm.after(resp, &mut context).await;
             }
 
-            match params.itsi_server_token_preference {
+            match token_preference {
                 ItsiServerTokenPreference::Version => {
                     resp.headers_mut().insert("Server", SERVER_TOKEN_VERSION);
                 }
@@ -238,28 +240,31 @@ impl Service<Request<Incoming>> for ItsiHttpService {
         };
 
         if let Some(timeout_duration) = request_timeout {
-            Box::pin(async move {
-                match timeout(timeout_duration, service_future).await {
-                    Ok(result) => result,
-                    Err(_) => {
-                        // If we're still running Ruby at this point, we can't just kill the
-                        // thread as it might be in a critical section.
-                        // Instead we must ask the worker to hot restart.
-                        if is_ruby_request.load(Ordering::Relaxed) {
-                            if is_single_mode {
-                                // If we're in single mode, re-exec the whole process
-                                send_lifecycle_event(LifecycleEvent::Restart);
-                            } else {
-                                // Otherwise we can shutdown the worker and rely on the master to restart it
-                                send_lifecycle_event(LifecycleEvent::Shutdown);
-                            }
+            match timeout(timeout_duration, service_future).await {
+                Ok(result) => result,
+                Err(_) => {
+                    // If we're still running Ruby at this point, we can't just kill the
+                    // thread as it might be in a critical section.
+                    // Instead we must ask the worker to hot restart.
+                    // But only if we're not already shutting down
+                    if is_ruby_request.load(Ordering::Relaxed)
+                        && !SHUTDOWN_REQUESTED.load(Ordering::SeqCst)
+                    {
+                        // When we've detected a timeout, use the safer send_lifecycle_event
+                        // which will properly handle signal-safe state transitions
+                        if is_single_mode {
+                            // If we're in single mode, re-exec the whole process
+                            send_lifecycle_event(LifecycleEvent::Restart);
+                        } else {
+                            // Otherwise we can shutdown the worker and rely on the master to restart it
+                            send_lifecycle_event(LifecycleEvent::Shutdown);
                         }
-                        Ok(TIMEOUT_RESPONSE.to_http_response(accept).await)
                     }
+                    Ok(TIMEOUT_RESPONSE.to_http_response(accept).await)
                 }
-            })
+            }
         } else {
-            Box::pin(service_future)
+            service_future.await
         }
     }
 }