itsi-server 0.1.1 → 0.1.3

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

@@ -0,0 +1,254 @@
+use crate::server::{
+    itsi_server::Server, lifecycle_event::LifecycleEvent, listener::Listener,
+    process_worker::ProcessWorker,
+};
+use itsi_error::{ItsiError, Result};
+use itsi_rb_helpers::{call_without_gvl, create_ruby_thread};
+use itsi_tracing::{error, info, warn};
+use nix::{
+    libc::{self, 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::instrument;
+pub(crate) struct ClusterMode {
+    pub listeners: Arc<Vec<Arc<Listener>>>,
+    pub server: Arc<Server>,
+    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: Arc<Server>,
+        listeners: Arc<Vec<Arc<Listener>>>,
+        lifecycle_channel: broadcast::Sender<LifecycleEvent>,
+    ) -> Self {
+        if let Some(f) = server.before_fork.lock().take() {
+            f();
+        }
+        let process_workers = (0..server.workers)
+            .map(|_| ProcessWorker {
+                worker_id: WORKER_ID.fetch_add(1, std::sync::atomic::Ordering::Relaxed),
+                ..Default::default()
+            })
+            .collect();
+
+        Self {
+            listeners,
+            server,
+            process_workers: parking_lot::Mutex::new(process_workers),
+            lifecycle_channel,
+        }
+    }
+
+    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")
+    }
+
+    #[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::Shutdown => {
+                self.shutdown().await?;
+                Ok(())
+            }
+            LifecycleEvent::Restart => {
+                for worker in self.process_workers.lock().iter() {
+                    worker.reboot(self.clone()).await?;
+                }
+                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.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();
+                }
+                unsafe { exit(0) };
+            }
+        }
+    }
+
+    pub async fn shutdown(&self) -> Result<()> {
+        let shutdown_timeout = self.server.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 => {
+              info!("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 fn receive_signal(signal: i32) {
+        match signal {
+            libc::SIGCHLD => {
+                CHILD_SIGNAL_SENDER.lock().as_ref().inspect(|i| {
+                    i.send(()).ok();
+                });
+            }
+            _ => {
+                // Handle other signals
+            }
+        }
+    }
+
+    pub fn stop(&self) -> Result<()> {
+        unsafe { libc::signal(libc::SIGCHLD, libc::SIG_DFL) };
+
+        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.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);
+
+        unsafe { libc::signal(libc::SIGCHLD, Self::receive_signal as usize) };
+
+        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 mut memory_check_interval = time::interval(time::Duration::from_secs(2));
+          loop {
+            tokio::select! {
+              _ = receiver.changed() => {
+                let mut workers = self_ref.process_workers.lock();
+                workers.retain(|worker| {
+                  worker.boot_if_dead(Arc::clone(&self_ref))
+                });
+                if workers.is_empty() {
+                    warn!("No workers running. Send SIGTTIN to increase worker count");
+                }
+              }
+              _ = memory_check_interval.tick() => {
+                if let Some(memory_limit) = self_ref.server.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();
+                      }
+                    }
+                  }
+                }
+              }
+              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),
+              }
+            }
+          }
+        });
+
+        Ok(())
+    }
+}

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

@@ -0,0 +1,27 @@
+use cluster_mode::ClusterMode;
+use itsi_error::Result;
+use single_mode::SingleMode;
+use std::sync::Arc;
+pub mod cluster_mode;
+pub mod single_mode;
+
+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.clone().run(),
+            ServeStrategy::Cluster(cluster_router) => cluster_router.clone().run(),
+        }
+    }
+
+    pub(crate) fn stop(&self) -> Result<()> {
+        match self {
+            ServeStrategy::Single(single_router) => single_router.clone().stop(),
+            ServeStrategy::Cluster(cluster_router) => cluster_router.clone().stop(),
+        }
+    }
+}

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

@@ -0,0 +1,241 @@
+use crate::{
+    request::itsi_request::ItsiRequest,
+    server::{
+        io_stream::IoStream,
+        itsi_server::{RequestJob, Server},
+        lifecycle_event::LifecycleEvent,
+        listener::{Listener, TokioListener},
+        thread_worker::{build_thread_workers, ThreadWorker},
+    },
+};
+use http::Request;
+use hyper::{body::Incoming, service::service_fn};
+use hyper_util::{
+    rt::{TokioExecutor, TokioIo, TokioTimer},
+    server::conn::auto::Builder,
+};
+use itsi_error::{ItsiError, Result};
+use itsi_tracing::{debug, error, info};
+use nix::unistd::Pid;
+use std::{
+    num::NonZeroU8,
+    pin::Pin,
+    sync::Arc,
+    time::{Duration, Instant},
+};
+use tokio::{
+    runtime::{Builder as RuntimeBuilder, Runtime},
+    sync::broadcast,
+    task::JoinSet,
+};
+use tracing::instrument;
+
+pub struct SingleMode {
+    pub executor: Builder<TokioExecutor>,
+    pub server: Arc<Server>,
+    pub sender: async_channel::Sender<RequestJob>,
+    pub(crate) listeners: Arc<Vec<Arc<Listener>>>,
+    pub(crate) thread_workers: Arc<Vec<ThreadWorker>>,
+    pub(crate) lifecycle_channel: broadcast::Sender<LifecycleEvent>,
+}
+
+pub enum RunningPhase {
+    Running,
+    ShutdownPending,
+    Shutdown,
+}
+
+impl SingleMode {
+    #[instrument(parent=None, skip_all, fields(pid=format!("{:?}", Pid::this())))]
+    pub(crate) fn new(
+        server: Arc<Server>,
+        listeners: Arc<Vec<Arc<Listener>>>,
+        lifecycle_channel: broadcast::Sender<LifecycleEvent>,
+    ) -> Result<Self> {
+        let (thread_workers, sender) = build_thread_workers(
+            Pid::this(),
+            NonZeroU8::try_from(server.threads).unwrap(),
+            server.app,
+            server.scheduler_class.clone(),
+        )?;
+        Ok(Self {
+            executor: Builder::new(TokioExecutor::new()),
+            listeners,
+            server,
+            sender,
+            thread_workers,
+            lifecycle_channel,
+        })
+    }
+
+    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 stop(&self) -> Result<()> {
+        Ok(())
+    }
+
+    #[instrument(parent=None, skip(self))]
+    pub fn run(self: Arc<Self>) -> Result<()> {
+        let mut listener_task_set = JoinSet::new();
+        let self_ref = Arc::new(self);
+        self_ref.build_runtime().block_on(async {
+
+          for listener in self_ref.listeners.clone().iter() {
+              let listener = Arc::new(listener.to_tokio_listener());
+              let mut lifecycle_rx = self_ref.lifecycle_channel.subscribe();
+              let self_ref = self_ref.clone();
+              let listener = listener.clone();
+              let (shutdown_sender, mut shutdown_receiver) = tokio::sync::watch::channel::<RunningPhase>(RunningPhase::Running);
+              listener_task_set.spawn(async move {
+                let strategy = self_ref.clone();
+                loop {
+                    tokio::select! {
+                        accept_result = listener.accept() => match accept_result {
+                          Ok(accept_result) => {
+                            if let Err(e) = strategy.serve_connection(accept_result, listener.clone(), shutdown_receiver.clone()).await {
+                              error!("Error in serve_connection {:?}", e)
+                            }
+                          },
+                          Err(e) => debug!("Listener.accept failed {:?}", e),
+                        },
+                        _ = shutdown_receiver.changed() => {
+                          break;
+                        }
+                        lifecycle_event = lifecycle_rx.recv() => match lifecycle_event{
+                          Ok(lifecycle_event) => {
+                            if let Err(e) = strategy.handle_lifecycle_event(lifecycle_event, shutdown_sender.clone()).await{
+                              match e {
+                                ItsiError::Break() => break,
+                                _ => error!("Error in handle_lifecycle_event {:?}", e)
+                              }
+                            }
+
+                          },
+                          Err(e) => error!("Error receiving lifecycle_event: {:?}", e),
+                      }
+                    }
+                }
+                if let Ok(listener) = Arc::try_unwrap(listener){
+                  listener.unbind();
+                }
+            });
+
+          }
+
+          while let Some(_res) = listener_task_set.join_next().await {}
+        });
+
+        Ok(())
+    }
+
+    pub(crate) async fn serve_connection(
+        &self,
+        stream: IoStream,
+        listener: Arc<TokioListener>,
+        shutdown_channel: tokio::sync::watch::Receiver<RunningPhase>,
+    ) -> Result<()> {
+        let sender_clone = self.sender.clone();
+        let addr = stream.addr();
+        let io: TokioIo<Pin<Box<IoStream>>> = TokioIo::new(Box::pin(stream));
+        let server = self.server.clone();
+        let executor = self.executor.clone();
+        let mut shutdown_channel_clone = shutdown_channel.clone();
+        tokio::spawn(async move {
+            let server = server.clone();
+            let mut executor = executor.clone();
+            let mut binding = executor.http1();
+            let shutdown_channel = shutdown_channel_clone.clone();
+            let mut serve = Box::pin(
+                binding
+                    .timer(TokioTimer::new())
+                    .header_read_timeout(Duration::from_secs(1))
+                    .serve_connection_with_upgrades(
+                        io,
+                        service_fn(move |hyper_request: Request<Incoming>| {
+                            ItsiRequest::process_request(
+                                hyper_request,
+                                sender_clone.clone(),
+                                server.clone(),
+                                listener.clone(),
+                                addr.clone(),
+                                shutdown_channel.clone(),
+                            )
+                        }),
+                    ),
+            );
+
+            tokio::select! {
+                // Await the connection finishing naturally.
+                res = &mut serve => {
+                    match res{
+                        Ok(()) => {
+                          debug!("Connection closed normally")
+                        },
+                        Err(res) => {
+                          debug!("Connection finished with error: {:?}", res)
+                        }
+                    }
+                    serve.as_mut().graceful_shutdown();
+                },
+                // A lifecycle event triggers shutdown.
+                _ = shutdown_channel_clone.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);
+                    }
+                }
+            }
+        });
+        Ok(())
+    }
+
+    pub async fn handle_lifecycle_event(
+        &self,
+        lifecycle_event: LifecycleEvent,
+        shutdown_sender: tokio::sync::watch::Sender<RunningPhase>,
+    ) -> Result<()> {
+        if let LifecycleEvent::Shutdown = lifecycle_event {
+            shutdown_sender
+                .send(RunningPhase::ShutdownPending)
+                .expect("Failed to send shutdown pending signal");
+            let deadline = Instant::now() + Duration::from_secs_f64(self.server.shutdown_timeout);
+            for worker in &*self.thread_workers {
+                worker.request_shutdown().await;
+            }
+            while Instant::now() < deadline {
+                tokio::time::sleep(Duration::from_millis(50)).await;
+                let alive_threads = self
+                    .thread_workers
+                    .iter()
+                    .filter(|worker| worker.poll_shutdown(deadline))
+                    .count();
+                if alive_threads == 0 {
+                    break;
+                }
+                tokio::time::sleep(Duration::from_millis(200)).await;
+            }
+
+            info!("Sending shutdown signal");
+            shutdown_sender
+                .send(RunningPhase::Shutdown)
+                .expect("Failed to send shutdown signal");
+            self.thread_workers.iter().for_each(|worker| {
+                worker.poll_shutdown(deadline);
+            });
+
+            return Err(ItsiError::Break());
+        }
+        Ok(())
+    }
+}

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

@@ -0,0 +1,70 @@
+use std::sync::{atomic::AtomicI8, LazyLock};
+
+use nix::libc::{self, sighandler_t};
+use tokio::sync::{self, broadcast};
+
+use super::lifecycle_event::LifecycleEvent;
+
+pub static SIGNAL_HANDLER_CHANNEL: LazyLock<(
+    broadcast::Sender<LifecycleEvent>,
+    broadcast::Receiver<LifecycleEvent>,
+)> = LazyLock::new(|| sync::broadcast::channel(5));
+
+pub static SIGINT_COUNT: AtomicI8 = AtomicI8::new(0);
+fn receive_signal(signum: i32, _: sighandler_t) {
+    SIGINT_COUNT.fetch_add(-1, std::sync::atomic::Ordering::SeqCst);
+    match signum {
+        libc::SIGTERM | libc::SIGINT => {
+            SIGINT_COUNT.fetch_add(2, std::sync::atomic::Ordering::SeqCst);
+            if SIGINT_COUNT.load(std::sync::atomic::Ordering::SeqCst) < 2 {
+                SIGNAL_HANDLER_CHANNEL.0.send(LifecycleEvent::Shutdown).ok();
+            } else {
+                // Not messing about. Force shutdown.
+                SIGNAL_HANDLER_CHANNEL
+                    .0
+                    .send(LifecycleEvent::ForceShutdown)
+                    .ok();
+            }
+        }
+        libc::SIGUSR1 => {
+            SIGNAL_HANDLER_CHANNEL.0.send(LifecycleEvent::Restart).ok();
+        }
+        libc::SIGTTIN => {
+            SIGNAL_HANDLER_CHANNEL
+                .0
+                .send(LifecycleEvent::IncreaseWorkers)
+                .ok();
+        }
+        libc::SIGTTOU => {
+            SIGNAL_HANDLER_CHANNEL
+                .0
+                .send(LifecycleEvent::DecreaseWorkers)
+                .ok();
+        }
+        _ => {}
+    }
+}
+
+pub fn reset_signal_handlers() -> bool {
+    SIGINT_COUNT.store(0, std::sync::atomic::Ordering::SeqCst);
+    unsafe {
+        libc::signal(libc::SIGTERM, receive_signal as usize);
+        libc::signal(libc::SIGINT, receive_signal as usize);
+        libc::signal(libc::SIGUSR1, receive_signal as usize);
+        libc::signal(libc::SIGUSR2, receive_signal as usize);
+        libc::signal(libc::SIGTTIN, receive_signal as usize);
+        libc::signal(libc::SIGTTOU, receive_signal as usize);
+    }
+    true
+}
+
+pub fn clear_signal_handlers() {
+    unsafe {
+        libc::signal(libc::SIGTERM, libc::SIG_DFL);
+        libc::signal(libc::SIGINT, libc::SIG_DFL);
+        libc::signal(libc::SIGUSR1, libc::SIG_DFL);
+        libc::signal(libc::SIGUSR2, libc::SIG_DFL);
+        libc::signal(libc::SIGTTIN, libc::SIG_DFL);
+        libc::signal(libc::SIGTTOU, libc::SIG_DFL);
+    }
+}