itsi-scheduler 0.1.5 → 0.1.11

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

@@ -2,24 +2,26 @@ 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 = "localhost:3000";
+static DEFAULT_BIND: &str = "http://localhost:3000";
 
 #[magnus::wrap(class = "Itsi::Server", free_immediately, size)]
 #[derive(Clone)]
@@ -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,31 +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(())
         })?;
-        if let Ok(listeners) = Arc::try_unwrap(listeners) {
-            listeners.into_iter().for_each(|listener| {
-                if let Ok(listener) = Arc::try_unwrap(listener) {
-                    listener.unbind()
-                };
-            });
-        }
         clear_signal_handlers();
+        self.strategy.write().take();
+        info!("Server stopped");
         Ok(())
     }
 }

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

@@ -1,6 +1,7 @@
 use super::bind::{Bind, BindAddress};
 use super::bind_protocol::BindProtocol;
 use super::io_stream::IoStream;
+use super::serve_strategy::single_mode::RunningPhase;
 use super::tls::ItsiTlsAcceptor;
 use itsi_error::{ItsiError, Result};
 use itsi_tracing::info;
@@ -11,6 +12,7 @@ use std::{os::unix::net::UnixListener, path::PathBuf};
 use tokio::net::TcpListener as TokioTcpListener;
 use tokio::net::UnixListener as TokioUnixListener;
 use tokio::net::{unix, TcpStream, UnixStream};
+use tokio::sync::watch::Receiver;
 use tokio_rustls::TlsAcceptor;
 use tokio_stream::StreamExt;
 use tracing::error;
@@ -23,45 +25,79 @@ pub(crate) enum Listener {
 }
 
 pub(crate) enum TokioListener {
-    Tcp {
-        listener: TokioTcpListener,
-        host: String,
-        port: u16,
-    },
-    TcpTls {
-        listener: TokioTcpListener,
-        acceptor: ItsiTlsAcceptor,
-        host: String,
-        port: u16,
-    },
-    Unix {
-        listener: TokioUnixListener,
-    },
-    UnixTls {
-        listener: TokioUnixListener,
-        acceptor: ItsiTlsAcceptor,
-    },
+    Tcp(TokioTcpListener),
+    TcpTls(TokioTcpListener, ItsiTlsAcceptor),
+    Unix(TokioUnixListener),
+    UnixTls(TokioUnixListener, ItsiTlsAcceptor),
+}
+
+#[derive(Debug, Clone)]
+pub struct ListenerInfo {
+    pub host: String,
+    pub port: u16,
+    pub scheme: String,
 }
 
 impl TokioListener {
-    pub fn unbind(self) {
+    pub fn listener_info(&self) -> ListenerInfo {
         match self {
-            TokioListener::Tcp { listener, .. } => drop(listener.into_std().unwrap()),
-            TokioListener::TcpTls { listener, .. } => drop(listener.into_std().unwrap()),
-            TokioListener::Unix { listener } => drop(listener.into_std().unwrap()),
-            TokioListener::UnixTls { listener, .. } => drop(listener.into_std().unwrap()),
-        };
+            TokioListener::Tcp(listener) => ListenerInfo {
+                host: listener
+                    .local_addr()
+                    .unwrap()
+                    .ip()
+                    .to_canonical()
+                    .to_string(),
+                port: listener.local_addr().unwrap().port(),
+                scheme: "http".to_string(),
+            },
+            TokioListener::TcpTls(listener, _) => ListenerInfo {
+                host: listener
+                    .local_addr()
+                    .unwrap()
+                    .ip()
+                    .to_canonical()
+                    .to_string(),
+                port: listener.local_addr().unwrap().port(),
+                scheme: "https".to_string(),
+            },
+            TokioListener::Unix(listener) => ListenerInfo {
+                host: listener
+                    .local_addr()
+                    .unwrap()
+                    .as_pathname()
+                    .unwrap()
+                    .to_str()
+                    .unwrap()
+                    .to_owned(),
+                port: 0,
+                scheme: "unix".to_string(),
+            },
+            TokioListener::UnixTls(listener, _) => ListenerInfo {
+                host: listener
+                    .local_addr()
+                    .unwrap()
+                    .as_pathname()
+                    .unwrap()
+                    .to_str()
+                    .unwrap()
+                    .to_owned(),
+                port: 0,
+                scheme: "ssl".to_string(),
+            },
+        }
     }
+
     pub(crate) async fn accept(&self) -> Result<IoStream> {
         match self {
-            TokioListener::Tcp { listener, .. } => TokioListener::accept_tcp(listener).await,
-            TokioListener::TcpTls {
-                listener, acceptor, ..
-            } => TokioListener::accept_tls(listener, acceptor).await,
-            TokioListener::Unix { listener, .. } => TokioListener::accept_unix(listener).await,
-            TokioListener::UnixTls {
-                listener, acceptor, ..
-            } => TokioListener::accept_unix_tls(listener, acceptor).await,
+            TokioListener::Tcp(listener) => TokioListener::accept_tcp(listener).await,
+            TokioListener::TcpTls(listener, acceptor) => {
+                TokioListener::accept_tls(listener, acceptor).await
+            }
+            TokioListener::Unix(listener) => TokioListener::accept_unix(listener).await,
+            TokioListener::UnixTls(listener, acceptor) => {
+                TokioListener::accept_unix_tls(listener, acceptor).await
+            }
         }
     }
 
@@ -70,17 +106,24 @@ impl TokioListener {
         Self::to_tokio_io(Stream::TcpStream(tcp_stream), None).await
     }
 
-    pub async fn spawn_state_task(&self) {
-        if let TokioListener::TcpTls {
-            acceptor: ItsiTlsAcceptor::Automatic(_acme_acceptor, state, _server_config),
-            ..
-        } = self
+    pub async fn spawn_state_task(&self, mut shutdown_receiver: Receiver<RunningPhase>) {
+        if let TokioListener::TcpTls(
+            _,
+            ItsiTlsAcceptor::Automatic(_acme_acceptor, state, _server_config),
+        ) = self
         {
             let mut state = state.lock().await;
             loop {
-                match StreamExt::next(&mut *state).await {
-                    Some(event) => info!("Received acme event: {:?}", event),
-                    None => error!("Received no acme event"),
+                tokio::select! {
+                  stream_event = StreamExt::next(&mut *state) => {
+                      match stream_event {
+                        Some(event) => info!("ACME Event: {:?}", event),
+                        None => error!("Received no acme event"),
+                      }
+                  },
+                  _ = shutdown_receiver.changed() => {
+                      break;
+                  }
                 }
             }
         }
@@ -175,33 +218,6 @@ impl TokioListener {
             },
         }
     }
-
-    pub(crate) fn scheme(&self) -> String {
-        match self {
-            TokioListener::Tcp { .. } => "http".to_string(),
-            TokioListener::TcpTls { .. } => "https".to_string(),
-            TokioListener::Unix { .. } => "http".to_string(),
-            TokioListener::UnixTls { .. } => "https".to_string(),
-        }
-    }
-
-    pub(crate) fn port(&self) -> u16 {
-        match self {
-            TokioListener::Tcp { port, .. } => *port,
-            TokioListener::TcpTls { port, .. } => *port,
-            TokioListener::Unix { .. } => 0,
-            TokioListener::UnixTls { .. } => 0,
-        }
-    }
-
-    pub(crate) fn host(&self) -> String {
-        match self {
-            TokioListener::Tcp { host, .. } => host.to_string(),
-            TokioListener::TcpTls { host, .. } => host.to_string(),
-            TokioListener::Unix { .. } => "unix".to_string(),
-            TokioListener::UnixTls { .. } => "unix".to_string(),
-        }
-    }
 }
 
 enum Stream {
@@ -227,48 +243,22 @@ impl std::fmt::Display for SockAddr {
 }
 
 impl Listener {
-    pub fn unbind(self) {
+    pub fn into_tokio_listener(self) -> TokioListener {
         match self {
-            Listener::Tcp(listener) => drop(listener),
-            Listener::TcpTls((listener, _)) => drop(listener),
-            Listener::Unix(listener) => drop(listener),
-            Listener::UnixTls((listener, _)) => drop(listener),
-        };
-    }
-    pub fn to_tokio_listener(&self) -> TokioListener {
-        match self {
-            Listener::Tcp(listener) => TokioListener::Tcp {
-                listener: TokioTcpListener::from_std(TcpListener::try_clone(listener).unwrap())
-                    .unwrap(),
-                host: listener
-                    .local_addr()
-                    .unwrap()
-                    .ip()
-                    .to_canonical()
-                    .to_string(),
-                port: listener.local_addr().unwrap().port(),
-            },
-            Listener::TcpTls((listener, acceptor)) => TokioListener::TcpTls {
-                listener: TokioTcpListener::from_std(TcpListener::try_clone(listener).unwrap())
-                    .unwrap(),
-                acceptor: acceptor.clone(),
-                host: listener
-                    .local_addr()
-                    .unwrap()
-                    .ip()
-                    .to_canonical()
-                    .to_string(),
-                port: listener.local_addr().unwrap().port(),
-            },
-            Listener::Unix(listener) => TokioListener::Unix {
-                listener: TokioUnixListener::from_std(UnixListener::try_clone(listener).unwrap())
-                    .unwrap(),
-            },
-            Listener::UnixTls((listener, acceptor)) => TokioListener::UnixTls {
-                listener: TokioUnixListener::from_std(UnixListener::try_clone(listener).unwrap())
-                    .unwrap(),
-                acceptor: acceptor.clone(),
-            },
+            Listener::Tcp(listener) => {
+                TokioListener::Tcp(TokioTcpListener::from_std(listener).unwrap())
+            }
+            Listener::TcpTls((listener, acceptor)) => TokioListener::TcpTls(
+                TokioTcpListener::from_std(listener).unwrap(),
+                acceptor.clone(),
+            ),
+            Listener::Unix(listener) => {
+                TokioListener::Unix(TokioUnixListener::from_std(listener).unwrap())
+            }
+            Listener::UnixTls((listener, acceptor)) => TokioListener::UnixTls(
+                TokioUnixListener::from_std(listener).unwrap(),
+                acceptor.clone(),
+            ),
         }
     }
 }
@@ -307,6 +297,7 @@ fn connect_tcp_socket(addr: IpAddr, port: u16) -> Result<TcpListener> {
     socket.set_nonblocking(true).ok();
     socket.set_nodelay(true).ok();
     socket.set_recv_buffer_size(262_144).ok();
+    info!("Binding to {:?}", socket_address);
     socket.bind(&socket_address.into())?;
     socket.listen(1024)?;
     Ok(socket.into())

data/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();