itsi-server 0.1.1 → 0.2.2

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(),
+        }
+    }
+}
+
+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()) {
+                    Ok(single_mode) => {
+                        if cluster_template
+                            .server_config
+                            .server_params
+                            .read()
+                            .pin_worker_cores
+                        {
+                            core_affinity::set_for_current(
+                                CORE_IDS[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() {
+            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() {
+                info!("Worker still alive, sending SIGKILL {}", pid);
+                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/cluster_mode.rs

@@ -0,0 +1,354 @@
+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_with_gvl, call_without_gvl, create_ruby_thread};
+use itsi_tracing::{error, info, warn};
+use magnus::Value;
+use nix::{libc::exit, unistd::Pid};
+
+use std::{
+    sync::{atomic::AtomicUsize, Arc},
+    time::{Duration, Instant},
+};
+use tokio::{
+    runtime::{Builder as RuntimeBuilder, Runtime},
+    sync::{broadcast, 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);
+
+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),
+                ..Default::default()
+            })
+            .collect();
+
+        Self {
+            server_config,
+            process_workers: parking_lot::Mutex::new(process_workers),
+            lifecycle_channel: SIGNAL_HANDLER_CHANNEL.0.clone(),
+        }
+    }
+
+    pub fn build_runtime(&self) -> Runtime {
+        let mut builder: RuntimeBuilder = RuntimeBuilder::new_current_thread();
+        builder
+            .thread_name("itsi-server-accept-loop")
+            .thread_stack_size(3 * 1024 * 1024)
+            .enable_io()
+            .enable_time()
+            .build()
+            .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>,
+        lifecycle_event: LifecycleEvent,
+    ) -> Result<()> {
+        match lifecycle_event {
+            LifecycleEvent::Start => Ok(()),
+            LifecycleEvent::PrintInfo => {
+                self.print_info().await?;
+                Ok(())
+            }
+            LifecycleEvent::Shutdown => {
+                self.server_config.stop_watcher()?;
+                self.shutdown().await?;
+                self.invoke_hook("before_shutdown");
+                Ok(())
+            }
+            LifecycleEvent::Restart => {
+                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()?;
+                    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);
+                    }
+                }
+                self.process_workers.lock().extend(next_workers);
+                while workers_to_load > 0 {
+                    let mut workers = self.process_workers.lock();
+                    let worker = ProcessWorker {
+                        worker_id: WORKER_ID.fetch_add(1, std::sync::atomic::Ordering::Relaxed),
+                        ..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();
+                        })
+                    });
+                    workers.push(worker);
+                    workers_to_load -= 1
+                }
+                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),
+                    ..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();
+                    })
+                });
+                workers.push(worker);
+                Ok(())
+            }
+            LifecycleEvent::DecreaseWorkers => {
+                let worker = {
+                    let mut workers = self.process_workers.lock();
+                    workers.pop()
+                };
+                if let Some(dropped_worker) = worker {
+                    dropped_worker.request_shutdown();
+                    let force_kill_time = Instant::now()
+                        + Duration::from_secs_f64(
+                            self.server_config.server_params.read().shutdown_timeout,
+                        );
+                    while dropped_worker.is_alive() && force_kill_time > Instant::now() {
+                        tokio::time::sleep(Duration::from_millis(100)).await;
+                    }
+                    if dropped_worker.is_alive() {
+                        dropped_worker.force_kill();
+                    }
+                };
+                Ok(())
+            }
+            LifecycleEvent::ForceShutdown => {
+                for worker in self.process_workers.lock().iter() {
+                    worker.force_kill();
+                }
+                error!("Force shutdown!");
+                unsafe { exit(0) };
+            }
+            LifecycleEvent::ChildTerminated => {
+                CHILD_SIGNAL_SENDER.lock().as_ref().inspect(|i| {
+                    i.send(()).ok();
+                });
+                Ok(())
+            }
+        }
+    }
+
+    pub async fn shutdown(&self) -> Result<()> {
+        let shutdown_timeout = self.server_config.server_params.read().shutdown_timeout;
+        let workers = self.process_workers.lock().clone();
+
+        workers.iter().for_each(|worker| worker.request_shutdown());
+
+        let remaining_children = Arc::new(Mutex::new(workers.len()));
+        let monitor_handle = {
+            let remaining_children: Arc<Mutex<usize>> = Arc::clone(&remaining_children);
+            let mut workers = workers.clone();
+            tokio::spawn(async move {
+                loop {
+                    // Check if all workers have exited
+                    let mut remaining = remaining_children.lock().await;
+                    workers.retain(|worker| worker.is_alive());
+                    *remaining = workers.len();
+                    if *remaining == 0 {
+                        break;
+                    }
+                    sleep(Duration::from_millis(100)).await;
+                }
+            })
+        };
+
+        tokio::select! {
+            _ = monitor_handle => {
+              debug!("All children exited early, exit normally")
+            }
+            _ = sleep(Duration::from_secs_f64(shutdown_timeout)) => {
+                warn!("Graceful shutdown timeout reached, force killing remaining children");
+                workers.iter().for_each(|worker| worker.force_kill());
+            }
+        }
+
+        Err(ItsiError::Break)
+    }
+
+    pub async fn print_info(self: Arc<Self>) -> Result<()> {
+        println!("Itsi Cluster Info:");
+        println!("Master PID: {:?}", Pid::this());
+        if let Some(memory_limit) = self.server_config.server_params.read().worker_memory_limit {
+            println!("Worker Memory Limit: {}", memory_limit);
+        }
+
+        if self.server_config.watcher_fd.is_some() {
+            println!("File Watcher Enabled: true",);
+            if let Some(watchers) = self
+                .server_config
+                .server_params
+                .read()
+                .notify_watchers
+                .as_ref()
+            {
+                for watcher in watchers {
+                    println!(
+                        "Watching path: {} => {}",
+                        watcher.0,
+                        watcher
+                            .1
+                            .iter()
+                            .map(|path| path.join(","))
+                            .collect::<Vec<String>>()
+                            .join(" ")
+                    );
+                }
+            }
+        }
+        println!(
+            "Silent Mode: {}",
+            self.server_config.server_params.read().silence
+        );
+        println!(
+            "Preload: {}",
+            self.server_config.server_params.read().preload
+        );
+        let workers = self.process_workers.lock().clone();
+        for worker in workers {
+            worker.print_info()?;
+            sleep(Duration::from_millis(50)).await;
+        }
+        Ok(())
+    }
+
+    pub fn stop(&self) -> Result<()> {
+        for worker in self.process_workers.lock().iter() {
+            if worker.is_alive() {
+                worker.force_kill();
+            }
+        }
+        Ok(())
+    }
+
+    #[instrument(skip(self), fields(mode = "cluster", pid=format!("{:?}", Pid::this())))]
+    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)))?;
+
+        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 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)
+          } 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() => {
+                let mut workers = self_ref.process_workers.lock();
+                workers.retain(|worker| {
+                  worker.boot_if_dead(self_ref.clone())
+                });
+                if workers.is_empty() {
+                    warn!("No workers running. Send SIGTTIN to increase worker count");
+                }
+              }
+              _ = memory_check_interval.tick() => {
+                let worker_memory_limit = self_ref.server_config.server_params.read().worker_memory_limit;
+                if let Some(memory_limit) = worker_memory_limit {
+                  let largest_worker = {
+                    let workers = self_ref.process_workers.lock();
+                    workers.iter().max_by(|wa, wb| wa.memory_usage().cmp(&wb.memory_usage())).cloned()
+                  };
+                  if let Some(largest_worker) = largest_worker {
+                    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_limit_reached") {
+                          call_with_gvl(|_|  hook.call::<_, Value>((largest_worker.pid(),)).ok() );
+                        }
+                      }
+                    }
+                  }
+                }
+              }
+              lifecycle_event = lifecycle_rx.recv() => match lifecycle_event{
+                Ok(lifecycle_event) => {
+                  if let Err(e) = self_ref.clone().handle_lifecycle_event(lifecycle_event).await{
+                    match e {
+                      ItsiError::Break => break,
+                      _ => error!("Error in handle_lifecycle_event {:?}", e)
+                    }
+                  }
+
+                },
+                Err(e) => error!("Error receiving lifecycle_event: {:?}", e),
+              }
+            }
+          }
+        });
+        self.server_config
+            .server_params
+            .write()
+            .listeners
+            .lock()
+            .drain(..);
+        Ok(())
+    }
+}

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

@@ -0,0 +1,30 @@
+use std::sync::Arc;
+
+use cluster_mode::ClusterMode;
+use itsi_error::Result;
+use single_mode::SingleMode;
+
+pub mod cluster_mode;
+pub mod single_mode;
+
+#[derive(Clone)]
+pub(crate) enum ServeStrategy {
+    Single(Arc<SingleMode>),
+    Cluster(Arc<ClusterMode>),
+}
+
+impl ServeStrategy {
+    pub fn run(self) -> Result<()> {
+        match self {
+            ServeStrategy::Single(single_router) => single_router.run(),
+            ServeStrategy::Cluster(cluster_router) => cluster_router.run(),
+        }
+    }
+
+    pub(crate) fn stop(&self) -> Result<()> {
+        match self {
+            ServeStrategy::Single(single_router) => single_router.stop(),
+            ServeStrategy::Cluster(cluster_router) => cluster_router.stop(),
+        }
+    }
+}