itsi-scheduler 0.2.22-aarch64-linux

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

@@ -0,0 +1,247 @@
+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;
+use nix::{
+    errno::Errno,
+    sys::{
+        signal::{
+            kill,
+            Signal::{SIGKILL, SIGTERM, SIGUSR2},
+        },
+        wait::{waitpid, WaitPidFlag, WaitStatus},
+    },
+    unistd::{setpgid, Pid},
+};
+use parking_lot::Mutex;
+use std::{
+    process::{self, exit},
+    sync::{Arc, LazyLock},
+    time::{Duration, Instant},
+};
+use sysinfo::System;
+
+use tokio::{sync::watch, time::sleep};
+use tracing::{info, instrument, warn};
+
+#[derive(Clone, Debug)]
+pub struct ProcessWorker {
+    pub worker_id: usize,
+    pub child_pid: Arc<Mutex<Option<Pid>>>,
+    pub started_at: Instant,
+}
+
+impl Default for ProcessWorker {
+    fn default() -> Self {
+        Self {
+            worker_id: 0,
+            child_pid: Arc::new(Mutex::new(None)),
+            started_at: Instant::now(),
+        }
+    }
+}
+
+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))]
+    pub(crate) fn boot(&self, cluster_template: Arc<ClusterMode>) -> Result<()> {
+        let child_pid = *self.child_pid.lock();
+        if let Some(pid) = child_pid {
+            if self.is_alive() {
+                if let Err(e) = kill(pid, SIGTERM) {
+                    info!("Failed to send SIGTERM to process {}: {}", pid, e);
+                }
+            }
+            *self.child_pid.lock() = None;
+        }
+
+        match call_with_gvl(|_ruby| {
+            fork(
+                cluster_template
+                    .server_config
+                    .server_params
+                    .read()
+                    .hooks
+                    .get("after_fork")
+                    .cloned(),
+            )
+        }) {
+            Some(pid) => {
+                *self.child_pid.lock() = Some(Pid::from_raw(pid));
+            }
+            None => {
+                if let Err(e) = setpgid(
+                    Pid::from_raw(process::id() as i32),
+                    Pid::from_raw(process::id() as i32),
+                ) {
+                    error!("Failed to set process group ID: {}", e);
+                }
+                match SingleMode::new(cluster_template.server_config.clone(), self.worker_id) {
+                    Ok(single_mode) => {
+                        if cluster_template
+                            .server_config
+                            .server_params
+                            .read()
+                            .pin_worker_cores
+                        {
+                            core_affinity::set_for_current(
+                                CORE_IDS[(2 * self.worker_id) % CORE_IDS.len()],
+                            );
+                        }
+                        Arc::new(single_mode).run().ok();
+                    }
+                    Err(e) => {
+                        error!("Failed to boot into worker mode: {}", e);
+                    }
+                }
+                exit(0)
+            }
+        }
+        Ok(())
+    }
+
+    pub fn pid(&self) -> i32 {
+        if let Some(pid) = *self.child_pid.lock() {
+            return pid.as_raw();
+        }
+        0
+    }
+
+    pub(crate) fn memory_usage(&self) -> Option<u64> {
+        if let Some(pid) = *self.child_pid.lock() {
+            let s = System::new_all();
+            if let Some(process) = s.process(sysinfo::Pid::from(pid.as_raw() as usize)) {
+                return Some(process.memory());
+            }
+        }
+        None
+    }
+
+    pub(crate) async fn reboot(&self, cluster_template: Arc<ClusterMode>) -> Result<bool> {
+        self.graceful_shutdown(cluster_template.clone()).await;
+        let self_clone = self.clone();
+        let (booted_sender, mut booted_receiver) = watch::channel(false);
+        create_ruby_thread(move || {
+            call_without_gvl(move || {
+                if self_clone.boot(cluster_template).is_ok() {
+                    booted_sender.send(true).ok()
+                } else {
+                    booted_sender.send(false).ok()
+                };
+            })
+        });
+
+        booted_receiver
+            .changed()
+            .await
+            .map_err(|_| ItsiError::InternalServerError("Failed to boot worker".to_owned()))?;
+
+        let guard = booted_receiver.borrow();
+        let result = guard.to_owned();
+        // Not very robust, we should check to see if the worker is actually listening before considering this successful.
+        sleep(Duration::from_secs(1)).await;
+        Ok(result)
+    }
+
+    pub(crate) async fn graceful_shutdown(&self, cluster_template: Arc<ClusterMode>) {
+        let self_clone = self.clone();
+        self_clone.request_shutdown();
+        let force_kill_time = Instant::now()
+            + Duration::from_secs_f64(
+                cluster_template
+                    .server_config
+                    .server_params
+                    .read()
+                    .shutdown_timeout,
+            );
+        while self_clone.is_alive() && force_kill_time > Instant::now() {
+            tokio::time::sleep(Duration::from_millis(100)).await;
+        }
+        if self_clone.is_alive() {
+            self_clone.force_kill();
+        }
+    }
+
+    pub(crate) fn boot_if_dead(&self, cluster_template: Arc<ClusterMode>) -> bool {
+        if !self.is_alive() && self.child_pid.lock().is_some() {
+            if self.just_started() {
+                error!(
+                    "Worker in crash loop {:?}. Refusing to restart",
+                    self.child_pid.lock()
+                );
+                return false;
+            } else {
+                let self_clone = self.clone();
+                create_ruby_thread(move || {
+                    call_without_gvl(move || {
+                        self_clone.boot(cluster_template).ok();
+                    })
+                });
+            }
+        }
+        true
+    }
+
+    pub(crate) fn request_shutdown(&self) {
+        let child_pid = *self.child_pid.lock();
+        if let Some(pid) = child_pid {
+            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");
+            }
+        }
+    }
+
+    pub(crate) fn force_kill(&self) {
+        let child_pid = *self.child_pid.lock();
+        if let Some(pid) = child_pid {
+            if self.is_alive() {
+                if let Err(e) = kill(pid, SIGKILL) {
+                    error!("Failed to force kill process {}: {}", pid, e);
+                }
+            }
+        }
+    }
+
+    pub fn print_info(&self) -> Result<()> {
+        let child_pid = *self.child_pid.lock();
+        if let Some(pid) = child_pid {
+            println!("Worker {:?}, PID: {:?}", self.worker_id, pid);
+            if let Err(e) = kill(pid, SIGUSR2) {
+                error!("Failed to send SIGUSR2 to process {}: {}", pid, e);
+            }
+        }
+
+        Ok(())
+    }
+
+    pub(crate) fn just_started(&self) -> bool {
+        let now = Instant::now();
+        now.duration_since(self.started_at).as_millis() < 2000
+    }
+
+    pub(crate) fn is_alive(&self) -> bool {
+        let child_pid = *self.child_pid.lock();
+        if let Some(pid) = child_pid {
+            match waitpid(pid, Some(WaitPidFlag::WNOHANG)) {
+                Ok(WaitStatus::Exited(_, _)) | Ok(WaitStatus::Signaled(_, _, _)) => {
+                    return false;
+                }
+                Ok(WaitStatus::StillAlive) | Ok(_) => {}
+                Err(_) => return false,
+            }
+            match kill(pid, None) {
+                Ok(_) => true,
+                Err(errno) => !matches!(errno, Errno::ESRCH),
+            }
+        } else {
+            false
+        }
+    }
+}

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

@@ -0,0 +1,26 @@
+use http::StatusCode;
+use serde::Deserialize;
+
+#[derive(Debug, Clone, Deserialize, Default)]
+pub enum RedirectType {
+    #[serde(rename(deserialize = "permanent"))]
+    #[default]
+    Permanent,
+    #[serde(rename(deserialize = "temporary"))]
+    Temporary,
+    #[serde(rename(deserialize = "found"))]
+    Found,
+    #[serde(rename(deserialize = "moved_permanently"))]
+    MovedPermanently,
+}
+
+impl RedirectType {
+    pub fn status_code(&self) -> StatusCode {
+        match self {
+            RedirectType::Permanent => StatusCode::PERMANENT_REDIRECT,
+            RedirectType::Temporary => StatusCode::TEMPORARY_REDIRECT,
+            RedirectType::Found => StatusCode::FOUND,
+            RedirectType::MovedPermanently => StatusCode::MOVED_PERMANENTLY,
+        }
+    }
+}

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

@@ -0,0 +1,11 @@
+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;
+
+#[derive(Debug)]
+pub enum RequestJob {
+    ProcessHttpRequest(ItsiHttpRequest, Arc<HeapValue<Proc>>),
+    ProcessGrpcRequest(ItsiGrpcCall, Arc<HeapValue<Proc>>),
+    Shutdown,
+}

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.");
+            }
+        }
+    }
+}