itsi 0.1.8 → 0.1.11

data/gems/scheduler/ext/itsi_server/src/server/itsi_server.rs

@@ -2,21 +2,23 @@ use super::{
     bind::Bind,
     listener::Listener,
     serve_strategy::{cluster_mode::ClusterMode, single_mode::SingleMode},
-    signal::{clear_signal_handlers, reset_signal_handlers, SIGNAL_HANDLER_CHANNEL},
+    signal::{
+        clear_signal_handlers, reset_signal_handlers, send_shutdown_event, SIGNAL_HANDLER_CHANNEL,
+    },
 };
 use crate::{request::itsi_request::ItsiRequest, server::serve_strategy::ServeStrategy};
 use derive_more::Debug;
-use itsi_rb_helpers::call_without_gvl;
-use itsi_tracing::error;
+use itsi_rb_helpers::{call_without_gvl, HeapVal, HeapValue};
+use itsi_tracing::{error, run_silently};
 use magnus::{
     block::Proc,
     error::Result,
-    scan_args::{get_kwargs, scan_args, Args, KwArgs},
-    value::{InnerValue, Opaque, ReprValue},
-    RHash, Ruby, Symbol, Value,
+    scan_args::{get_kwargs, scan_args, Args, KwArgs, ScanArgsKw, ScanArgsOpt, ScanArgsRequired},
+    value::ReprValue,
+    ArgList, RArray, RHash, Ruby, Symbol, Value,
 };
-use parking_lot::Mutex;
-use std::{cmp::max, ops::Deref, sync::Arc};
+use parking_lot::{Mutex, RwLock};
+use std::{cmp::max, collections::HashMap, ops::Deref, sync::Arc};
 use tracing::{info, instrument};
 
 static DEFAULT_BIND: &str = "http://localhost:3000";
@@ -34,12 +36,11 @@ impl Deref for Server {
         &self.config
     }
 }
-type AfterFork = Mutex<Arc<Option<Box<dyn Fn() + Send + Sync>>>>;
 
 #[derive(Debug)]
 pub struct ServerConfig {
     #[debug(skip)]
-    pub app: Opaque<Value>,
+    pub app: HeapVal,
     #[allow(unused)]
     pub workers: u8,
     #[allow(unused)]
@@ -49,12 +50,14 @@ pub struct ServerConfig {
     pub script_name: String,
     pub(crate) binds: Mutex<Vec<Bind>>,
     #[debug(skip)]
-    pub before_fork: Mutex<Option<Box<dyn FnOnce() + Send + Sync>>>,
-    #[debug(skip)]
-    pub after_fork: AfterFork,
+    pub hooks: HashMap<String, HeapValue<Proc>>,
     pub scheduler_class: Option<String>,
     pub stream_body: Option<bool>,
     pub worker_memory_limit: Option<u64>,
+    #[debug(skip)]
+    pub(crate) strategy: RwLock<Option<ServeStrategy>>,
+    pub silence: bool,
+    pub oob_gc_responses_threshold: Option<u64>,
 }
 
 #[derive(Debug)]
@@ -63,6 +66,30 @@ pub enum RequestJob {
     Shutdown,
 }
 
+fn extract_args<Req, Opt, Splat>(
+    scan_args: &Args<(), (), (), (), RHash, ()>,
+    primaries: &[&str],
+    rest: &[&str],
+) -> Result<KwArgs<Req, Opt, Splat>>
+where
+    Req: ScanArgsRequired,
+    Opt: ScanArgsOpt,
+    Splat: ScanArgsKw,
+{
+    let symbols: Vec<Symbol> = primaries
+        .iter()
+        .chain(rest.iter())
+        .map(|&name| Symbol::new(name))
+        .collect();
+
+    let hash = scan_args
+        .keywords
+        .funcall::<_, _, RHash>("slice", symbols.into_arg_list_with(&Ruby::get().unwrap()))
+        .unwrap();
+
+    get_kwargs(hash, primaries, rest)
+}
+
 impl Server {
     #[instrument(
         name = "Itsi",
@@ -73,39 +100,44 @@ impl Server {
     pub fn new(args: &[Value]) -> Result<Self> {
         let scan_args: Args<(), (), (), (), RHash, ()> = scan_args(args)?;
 
-        type ArgSet1 = (
-            Option<u8>,
-            Option<u8>,
-            Option<f64>,
-            Option<String>,
-            Option<Vec<String>>,
-            Option<Proc>,
-            Option<Proc>,
-            Option<String>,
-            Option<bool>,
-        );
-
-        type ArgSet2 = (Option<u64>,);
-
-        let args1: KwArgs<(Value,), ArgSet1, ()> = get_kwargs(
-            scan_args
-                .keywords
-                .funcall::<_, _, RHash>(
-                    "slice",
-                    (
-                        Symbol::new("app"),
-                        Symbol::new("workers"),
-                        Symbol::new("threads"),
-                        Symbol::new("shutdown_timeout"),
-                        Symbol::new("script_name"),
-                        Symbol::new("binds"),
-                        Symbol::new("before_fork"),
-                        Symbol::new("after_fork"),
-                        Symbol::new("scheduler_class"),
-                        Symbol::new("stream_body"),
-                    ),
-                )
-                .unwrap(),
+        type Args1 = KwArgs<
+            (Value,),
+            (
+                // Workers
+                Option<u8>,
+                // Threads
+                Option<u8>,
+                // Shutdown Timeout
+                Option<f64>,
+                // Script Name
+                Option<String>,
+                // Binds
+                Option<Vec<String>>,
+                // Stream Body
+                Option<bool>,
+            ),
+            (),
+        >;
+
+        type Args2 = KwArgs<
+            (),
+            (
+                // Hooks
+                Option<RHash>,
+                // Scheduler Class
+                Option<String>,
+                // Worker Memory Limit
+                Option<u64>,
+                // Out-of-band GC Responses Threshold
+                Option<u64>,
+                // Silence
+                Option<bool>,
+            ),
+            (),
+        >;
+
+        let args1: Args1 = extract_args(
+            &scan_args,
             &["app"],
             &[
                 "workers",
@@ -113,24 +145,43 @@ impl Server {
                 "shutdown_timeout",
                 "script_name",
                 "binds",
-                "before_fork",
-                "after_fork",
-                "scheduler_class",
                 "stream_body",
             ],
         )?;
 
-        let args2: KwArgs<(), ArgSet2, ()> = get_kwargs(
-            scan_args
-                .keywords
-                .funcall::<_, _, RHash>("slice", (Symbol::new("worker_memory_limit"),))
-                .unwrap(),
+        let args2: Args2 = extract_args(
+            &scan_args,
             &[],
-            &["worker_memory_limit"],
+            &[
+                "hooks",
+                "scheduler_class",
+                "worker_memory_limit",
+                "oob_gc_responses_threshold",
+                "silence",
+            ],
         )?;
 
+        let hooks = args2
+            .optional
+            .0
+            .map(|rhash| -> Result<HashMap<String, HeapValue<Proc>>> {
+                let mut hook_map: HashMap<String, HeapValue<Proc>> = HashMap::new();
+                for pair in rhash.enumeratorize::<_, ()>("each", ()) {
+                    if let Some(pair_value) = RArray::from_value(pair?) {
+                        if let (Ok(key), Ok(value)) =
+                            (pair_value.entry::<Value>(0), pair_value.entry::<Proc>(1))
+                        {
+                            hook_map.insert(key.to_string(), HeapValue::from(value));
+                        }
+                    }
+                }
+                Ok(hook_map)
+            })
+            .transpose()?
+            .unwrap_or_default();
+
         let config = ServerConfig {
-            app: Opaque::from(args1.required.0),
+            app: HeapVal::from(args1.required.0),
             workers: max(args1.optional.0.unwrap_or(1), 1),
             threads: max(args1.optional.1.unwrap_or(1), 1),
             shutdown_timeout: args1.optional.2.unwrap_or(5.0),
@@ -144,33 +195,21 @@ impl Server {
                     .map(|s| s.parse())
                     .collect::<itsi_error::Result<Vec<Bind>>>()?,
             ),
-            before_fork: Mutex::new(args1.optional.5.map(|p| {
-                let opaque_proc = Opaque::from(p);
-                Box::new(move || {
-                    opaque_proc
-                        .get_inner_with(&Ruby::get().unwrap())
-                        .call::<_, Value>(())
-                        .unwrap();
-                }) as Box<dyn FnOnce() + Send + Sync>
-            })),
-            after_fork: Mutex::new(Arc::new(args1.optional.6.map(|p| {
-                let opaque_proc = Opaque::from(p);
-                Box::new(move || {
-                    opaque_proc
-                        .get_inner_with(&Ruby::get().unwrap())
-                        .call::<_, Value>(())
-                        .unwrap();
-                }) as Box<dyn Fn() + Send + Sync>
-            }))),
-            scheduler_class: args1.optional.7.clone(),
-            stream_body: args1.optional.8,
-            worker_memory_limit: args2.optional.0,
+            stream_body: args1.optional.5,
+            hooks,
+            scheduler_class: args2.optional.1.clone(),
+            worker_memory_limit: args2.optional.2,
+            strategy: RwLock::new(None),
+            oob_gc_responses_threshold: args2.optional.3,
+            silence: args2.optional.4.is_some_and(|s| s),
         };
 
-        if let Some(scheduler_class) = args1.optional.7 {
-            info!(scheduler_class, fiber_scheduler = true);
-        } else {
-            info!(fiber_scheduler = false);
+        if !config.silence {
+            if let Some(scheduler_class) = args2.optional.1 {
+                info!(scheduler_class, fiber_scheduler = true);
+            } else {
+                info!(fiber_scheduler = false);
+            }
         }
 
         Ok(Server {
@@ -179,7 +218,7 @@ impl Server {
     }
 
     #[instrument(name = "Bind", skip_all, fields(binds=format!("{:?}", self.config.binds.lock())))]
-    pub(crate) fn listeners(&self) -> Result<Arc<Vec<Arc<Listener>>>> {
+    pub(crate) fn build_listeners(&self) -> Result<Vec<Listener>> {
         let listeners = self
             .config
             .binds
@@ -189,17 +228,15 @@ impl Server {
             .map(Listener::try_from)
             .collect::<std::result::Result<Vec<Listener>, _>>()?
             .into_iter()
-            .map(Arc::new)
             .collect::<Vec<_>>();
         info!("Bound {:?} listeners", listeners.len());
-        Ok(Arc::new(listeners))
+        Ok(listeners)
     }
 
-    pub(crate) fn build_strategy(
-        self,
-        listeners: Arc<Vec<Arc<Listener>>>,
-    ) -> Result<ServeStrategy> {
+    pub(crate) fn build_strategy(self) -> Result<()> {
+        let listeners = self.build_listeners()?;
         let server = Arc::new(self);
+        let server_clone = server.clone();
 
         let strategy = if server.config.workers == 1 {
             ServeStrategy::Single(Arc::new(SingleMode::new(
@@ -214,24 +251,38 @@ impl Server {
                 SIGNAL_HANDLER_CHANNEL.0.clone(),
             )))
         };
-        Ok(strategy)
+
+        *server_clone.strategy.write() = Some(strategy);
+        Ok(())
+    }
+
+    pub fn stop(&self) -> Result<()> {
+        send_shutdown_event();
+        Ok(())
     }
 
     pub fn start(&self) -> Result<()> {
+        if self.silence {
+            run_silently(|| self.build_and_run_strategy())
+        } else {
+            self.build_and_run_strategy()
+        }
+    }
+
+    fn build_and_run_strategy(&self) -> Result<()> {
         reset_signal_handlers();
         let rself = self.clone();
-        let listeners = self.listeners()?;
-        let listeners_clone = listeners.clone();
         call_without_gvl(move || -> Result<()> {
-            let strategy = rself.build_strategy(listeners_clone)?;
-            if let Err(e) = strategy.run() {
+            rself.clone().build_strategy()?;
+            if let Err(e) = rself.strategy.read().as_ref().unwrap().run() {
                 error!("Error running server: {}", e);
-                strategy.stop()?;
+                rself.strategy.read().as_ref().unwrap().stop()?;
             }
-            drop(strategy);
             Ok(())
         })?;
         clear_signal_handlers();
+        self.strategy.write().take();
+        info!("Server stopped");
         Ok(())
     }
 }

data/gems/scheduler/ext/itsi_server/src/server/listener.rs

@@ -117,7 +117,7 @@ impl TokioListener {
                 tokio::select! {
                   stream_event = StreamExt::next(&mut *state) => {
                       match stream_event {
-                        Some(event) => info!("Received acme event: {:?}", event),
+                        Some(event) => info!("ACME Event: {:?}", event),
                         None => error!("Received no acme event"),
                       }
                   },
@@ -243,20 +243,20 @@ impl std::fmt::Display for SockAddr {
 }
 
 impl Listener {
-    pub fn to_tokio_listener(&self) -> TokioListener {
+    pub fn into_tokio_listener(self) -> TokioListener {
         match self {
-            Listener::Tcp(listener) => TokioListener::Tcp(
-                TokioTcpListener::from_std(TcpListener::try_clone(listener).unwrap()).unwrap(),
-            ),
+            Listener::Tcp(listener) => {
+                TokioListener::Tcp(TokioTcpListener::from_std(listener).unwrap())
+            }
             Listener::TcpTls((listener, acceptor)) => TokioListener::TcpTls(
-                TokioTcpListener::from_std(TcpListener::try_clone(listener).unwrap()).unwrap(),
+                TokioTcpListener::from_std(listener).unwrap(),
                 acceptor.clone(),
             ),
-            Listener::Unix(listener) => TokioListener::Unix(
-                TokioUnixListener::from_std(UnixListener::try_clone(listener).unwrap()).unwrap(),
-            ),
+            Listener::Unix(listener) => {
+                TokioListener::Unix(TokioUnixListener::from_std(listener).unwrap())
+            }
             Listener::UnixTls((listener, acceptor)) => TokioListener::UnixTls(
-                TokioUnixListener::from_std(UnixListener::try_clone(listener).unwrap()).unwrap(),
+                TokioUnixListener::from_std(listener).unwrap(),
                 acceptor.clone(),
             ),
         }

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

@@ -53,8 +53,8 @@ impl ProcessWorker {
             }
             *self.child_pid.lock() = None;
         }
-
-        match call_with_gvl(|_ruby| fork(cluster_template.server.after_fork.lock().clone())) {
+        match call_with_gvl(|_ruby| fork(cluster_template.server.hooks.get("after_fork").cloned()))
+        {
             Some(pid) => {
                 *self.child_pid.lock() = Some(Pid::from_raw(pid));
             }
@@ -67,7 +67,7 @@ impl ProcessWorker {
                 }
                 match SingleMode::new(
                     cluster_template.server.clone(),
-                    cluster_template.listeners.clone(),
+                    cluster_template.listeners.lock().drain(..).collect(),
                     cluster_template.lifecycle_channel.clone(),
                 ) {
                     Ok(single_mode) => {
@@ -83,6 +83,13 @@ impl ProcessWorker {
         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();

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

@@ -3,8 +3,11 @@ use crate::server::{
     process_worker::ProcessWorker,
 };
 use itsi_error::{ItsiError, Result};
-use itsi_rb_helpers::{call_without_gvl, create_ruby_thread};
+use itsi_rb_helpers::{
+    call_proc_and_log_errors, call_with_gvl, call_without_gvl, create_ruby_thread,
+};
 use itsi_tracing::{error, info, warn};
+use magnus::Value;
 use nix::{
     libc::{self, exit},
     unistd::Pid,
@@ -19,9 +22,9 @@ use tokio::{
     sync::{broadcast, watch, Mutex},
     time::{self, sleep},
 };
-use tracing::instrument;
+use tracing::{debug, instrument};
 pub(crate) struct ClusterMode {
-    pub listeners: Arc<Vec<Arc<Listener>>>,
+    pub listeners: parking_lot::Mutex<Vec<Listener>>,
     pub server: Arc<Server>,
     pub process_workers: parking_lot::Mutex<Vec<ProcessWorker>>,
     pub lifecycle_channel: broadcast::Sender<LifecycleEvent>,
@@ -34,12 +37,9 @@ static CHILD_SIGNAL_SENDER: parking_lot::Mutex<Option<watch::Sender<()>>> =
 impl ClusterMode {
     pub fn new(
         server: Arc<Server>,
-        listeners: Arc<Vec<Arc<Listener>>>,
+        listeners: Vec<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),
@@ -48,7 +48,7 @@ impl ClusterMode {
             .collect();
 
         Self {
-            listeners,
+            listeners: parking_lot::Mutex::new(listeners),
             server,
             process_workers: parking_lot::Mutex::new(process_workers),
             lifecycle_channel,
@@ -152,7 +152,7 @@ impl ClusterMode {
 
         tokio::select! {
             _ = monitor_handle => {
-              info!("All children exited early, exit normally")
+              debug!("All children exited early, exit normally")
             }
             _ = sleep(Duration::from_secs_f64(shutdown_timeout)) => {
                 warn!("Graceful shutdown timeout reached, force killing remaining children");
@@ -191,6 +191,9 @@ impl ClusterMode {
     #[instrument(skip(self), fields(mode = "cluster", pid=format!("{:?}", Pid::this())))]
     pub fn run(self: Arc<Self>) -> Result<()> {
         info!("Starting in Cluster mode");
+        if let Some(proc) = self.server.hooks.get("before_fork") {
+            call_with_gvl(|_| call_proc_and_log_errors(proc.clone()))
+        }
         self.process_workers
             .lock()
             .iter()
@@ -228,6 +231,9 @@ impl ClusterMode {
                     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.hooks.get("after_memory_threshold_reached") {
+                          call_with_gvl(|_|  hook.call::<_, Value>((largest_worker.pid(),)).ok() );
+                        }
                       }
                     }
                   }