itsi-server 0.1.1 → 0.1.9

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

@@ -1,56 +1,150 @@
 use super::{
     bind::Bind,
-    listener::{Listener, SockAddr},
+    listener::Listener,
+    serve_strategy::{cluster_mode::ClusterMode, single_mode::SingleMode},
+    signal::{
+        clear_signal_handlers, reset_signal_handlers, send_shutdown_event, SIGNAL_HANDLER_CHANNEL,
+    },
 };
-use crate::{request::itsi_request::ItsiRequest, ITSI_SERVER};
-use bytes::Bytes;
+use crate::{request::itsi_request::ItsiRequest, server::serve_strategy::ServeStrategy};
 use derive_more::Debug;
-use http_body_util::{combinators::BoxBody, Empty};
-use hyper::{
-    body::Incoming, header::HeaderName, service::service_fn, HeaderMap, Request, Response,
-    StatusCode,
-};
-use hyper_util::{rt::TokioExecutor, server::conn::auto::Builder};
-use itsi_tracing::{error, info};
+use itsi_rb_helpers::{call_without_gvl, HeapVal};
+use itsi_tracing::{error, run_silently};
 use magnus::{
+    block::Proc,
     error::Result,
-    scan_args::{get_kwargs, scan_args, Args, KwArgs},
-    value::{Opaque, ReprValue},
-    RHash, Ruby, Value,
+    scan_args::{get_kwargs, scan_args, Args, KwArgs, ScanArgsKw, ScanArgsOpt, ScanArgsRequired},
+    value::{InnerValue, Opaque, ReprValue},
+    ArgList, RHash, Ruby, Symbol, Value,
 };
-use parking_lot::Mutex;
-use std::{collections::HashMap, convert::Infallible, sync::Arc};
-use tokio::runtime::Builder as RuntimeBuilder;
-use tokio::task::JoinSet;
+use parking_lot::{Mutex, RwLock};
+use std::{cmp::max, ops::Deref, sync::Arc};
+use tracing::{info, instrument};
+
+static DEFAULT_BIND: &str = "http://localhost:3000";
 
 #[magnus::wrap(class = "Itsi::Server", free_immediately, size)]
-#[derive(Debug)]
+#[derive(Clone)]
 pub struct Server {
+    pub config: Arc<ServerConfig>,
+}
+
+impl Deref for Server {
+    type Target = ServerConfig;
+
+    fn deref(&self) -> &Self::Target {
+        &self.config
+    }
+}
+type AfterFork = Mutex<Arc<Option<Box<dyn Fn() + Send + Sync>>>>;
+
+#[derive(Debug)]
+pub struct ServerConfig {
     #[debug(skip)]
-    app: Opaque<Value>,
+    pub app: HeapVal,
     #[allow(unused)]
-    workers: u16,
+    pub workers: u8,
     #[allow(unused)]
-    threads: u16,
+    pub threads: u8,
     #[allow(unused)]
-    shutdown_timeout: f64,
-    script_name: String,
+    pub shutdown_timeout: f64,
+    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 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,
+}
+
+#[derive(Debug)]
+pub enum RequestJob {
+    ProcessRequest(ItsiRequest),
+    Shutdown,
+}
+
+// Define your helper function.
+// Here P, A, C correspond to the types for the first tuple, second tuple, and extra parameters respectively.
+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,
+{
+    // Combine the primary and rest names into one Vec of Symbols.
+    let symbols: Vec<Symbol> = primaries
+        .iter()
+        .chain(rest.iter())
+        .map(|&name| Symbol::new(name))
+        .collect();
+
+    // Call the "slice" function with the combined symbols.
+    let hash = scan_args
+        .keywords
+        .funcall::<_, _, RHash>("slice", symbols.into_arg_list_with(&Ruby::get().unwrap()))
+        .unwrap();
+
+    // Finally, call get_kwargs with the original name slices.
+    get_kwargs(hash, primaries, rest)
 }
 
 impl Server {
+    #[instrument(
+        name = "Itsi",
+        parent=None,
+        skip(args),
+        fields(workers = 1, threads = 1, shutdown_timeout = 5)
+    )]
     pub fn new(args: &[Value]) -> Result<Self> {
-        type OptionalArgs = (
-            Option<u16>,
-            Option<u16>,
-            Option<f64>,
-            Option<String>,
-            Option<Vec<String>>,
-        );
-
         let scan_args: Args<(), (), (), (), RHash, ()> = scan_args(args)?;
-        let args: KwArgs<(Value,), OptionalArgs, ()> = get_kwargs(
-            scan_args.keywords,
+
+        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<
+            (),
+            (
+                // Before Fork
+                Option<Proc>,
+                // After Fork
+                Option<Proc>,
+                // Scheduler Class
+                Option<String>,
+                // Worker Memory Limit
+                Option<u64>,
+                // Silence
+                Option<bool>,
+            ),
+            (),
+        >;
+
+        let args1: Args1 = extract_args(
+            &scan_args,
             &["app"],
             &[
                 "workers",
@@ -58,125 +152,143 @@ impl Server {
                 "shutdown_timeout",
                 "script_name",
                 "binds",
+                "stream_body",
+            ],
+        )?;
+
+        let args2: Args2 = extract_args(
+            &scan_args,
+            &[],
+            &[
+                "before_fork",
+                "after_fork",
+                "scheduler_class",
+                "worker_memory_limit",
+                "silence",
             ],
         )?;
-        let server = Server {
-            app: Opaque::from(args.required.0),
-            workers: args.optional.0.unwrap_or(1),
-            threads: args.optional.1.unwrap_or(1),
-            shutdown_timeout: args.optional.2.unwrap_or(5.0),
-            script_name: args.optional.3.unwrap_or("".to_string()),
+
+        let config = ServerConfig {
+            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),
+            script_name: args1.optional.3.unwrap_or("".to_string()),
             binds: Mutex::new(
-                args.optional
+                args1
+                    .optional
                     .4
-                    .unwrap_or_else(|| vec!["localhost:3000".to_string()])
+                    .unwrap_or_else(|| vec![DEFAULT_BIND.to_string()])
                     .into_iter()
-                    .map(|s| s.parse().unwrap_or_else(|_| Bind::default()))
-                    .collect(),
+                    .map(|s| s.parse())
+                    .collect::<itsi_error::Result<Vec<Bind>>>()?,
             ),
+            stream_body: args1.optional.5,
+            before_fork: Mutex::new(args2.optional.0.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(args2.optional.1.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: args2.optional.2.clone(),
+            worker_memory_limit: args2.optional.3,
+            silence: args2.optional.4.is_some_and(|s| s),
+            strategy: RwLock::new(None),
         };
-        Ok(server)
-    }
 
-    pub(crate) async fn process_request(
-        hyper_request: Request<Incoming>,
-        app: Opaque<Value>,
-        script_name: String,
-        listener: Arc<Listener>,
-        addr: SockAddr,
-    ) -> itsi_error::Result<Response<BoxBody<Bytes, Infallible>>> {
-        let request = ItsiRequest::build_from(hyper_request, addr, script_name, listener).await;
-        let ruby = Ruby::get().unwrap();
-        let server = ruby.get_inner(&ITSI_SERVER);
-        let response: Result<(u16, HashMap<String, String>, Value)> =
-            server.funcall("call", (app, request));
-        if let Ok((status, headers_raw, body)) = response {
-            let mut body_buf = vec![];
-            for body_chunk in body.enumeratorize("each", ()) {
-                body_buf.push(body_chunk.unwrap().to_string())
+        if !config.silence {
+            if let Some(scheduler_class) = args2.optional.2 {
+                info!(scheduler_class, fiber_scheduler = true);
+            } else {
+                info!(fiber_scheduler = false);
             }
-            body.check_funcall::<_, _, Value>("close", ());
-            let boxed_body = BoxBody::new(body_buf.join(""));
-            let mut response = Response::new(boxed_body);
-            let mut headers = HeaderMap::new();
-            headers_raw.into_iter().for_each(|(key, value)| {
-                let header_name: HeaderName = key.parse().unwrap();
-                headers.insert(header_name, value.parse().unwrap());
-            });
-            *response.headers_mut() = headers;
-            *response.status_mut() = StatusCode::from_u16(status).unwrap();
-            Ok(response)
+        }
+
+        Ok(Server {
+            config: Arc::new(config),
+        })
+    }
+
+    #[instrument(name = "Bind", skip_all, fields(binds=format!("{:?}", self.config.binds.lock())))]
+    pub(crate) fn build_listeners(&self) -> Result<Arc<Vec<Arc<Listener>>>> {
+        let listeners = self
+            .config
+            .binds
+            .lock()
+            .iter()
+            .cloned()
+            .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))
+    }
+
+    pub(crate) fn build_strategy(self, listeners: Arc<Vec<Arc<Listener>>>) -> Result<()> {
+        let server = Arc::new(self);
+        let server_clone = server.clone();
+
+        let strategy = if server.config.workers == 1 {
+            ServeStrategy::Single(Arc::new(SingleMode::new(
+                server,
+                listeners,
+                SIGNAL_HANDLER_CHANNEL.0.clone(),
+            )?))
+        } else {
+            ServeStrategy::Cluster(Arc::new(ClusterMode::new(
+                server,
+                listeners,
+                SIGNAL_HANDLER_CHANNEL.0.clone(),
+            )))
+        };
+
+        *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 {
-            let mut response = Response::new(BoxBody::new(Empty::new()));
-            *response.status_mut() = StatusCode::BAD_REQUEST;
-            Ok(response)
+            self.build_and_run_strategy()
         }
     }
 
-    pub fn start(&self) {
-        let mut builder: RuntimeBuilder = RuntimeBuilder::new_current_thread();
-        let runtime = builder
-            .thread_name("itsi-server-accept-loop")
-            .thread_stack_size(3 * 1024 * 1024)
-            .enable_io()
-            .enable_time()
-            .build()
-            .expect("Failed to build Tokio runtime");
-
-        runtime.block_on(async {
-            let server = Arc::new(Builder::new(TokioExecutor::new()));
-            let listeners: Vec<Listener> = self
-                .binds
-                .lock()
-                .iter()
-                .cloned()
-                .map(Listener::from)
-                .collect::<Vec<_>>();
-
-            let mut set = JoinSet::new();
-
-            for listener in listeners {
-                let app = self.app;
-                let server_clone = server.clone();
-                let listener_clone = Arc::new(listener);
-                let script_name = self.script_name.clone();
-
-                set.spawn(async move {
-                    loop {
-                        let server = server_clone.clone();
-                        let listener = listener_clone.clone();
-                        let script_name = script_name.clone();
-                        let (stream, addr) = match listener.accept().await {
-                            Ok(stream) => stream,
-                            Err(e) => {
-                                error!("Failed to accept connection: {:?}", e);
-                                continue;
-                            }
-                        };
-
-                        tokio::spawn(async move {
-                            if let Err(e) = server
-                                .serve_connection_with_upgrades(
-                                    stream,
-                                    service_fn(move |hyper_request: Request<Incoming>| {
-                                        Server::process_request(
-                                            hyper_request,
-                                            app,
-                                            script_name.clone(),
-                                            listener.clone(),
-                                            addr.clone(),
-                                        )
-                                    }),
-                                )
-                                .await
-                            {
-                                info!("Closed connection due to: {:?}", e);
-                            }
-                        });
-                    }
-                });
+    fn build_and_run_strategy(&self) -> Result<()> {
+        reset_signal_handlers();
+        let rself = self.clone();
+        let listeners = self.build_listeners()?;
+        let listeners_clone = listeners.clone();
+        call_without_gvl(move || -> Result<()> {
+            rself.clone().build_strategy(listeners_clone)?;
+            if let Err(e) = rself.clone().strategy.read().as_ref().unwrap().run() {
+                error!("Error running server: {}", e);
+                rself.strategy.read().as_ref().unwrap().stop()?;
             }
-            while let Some(_res) = set.join_next().await {}
-        })
+            Ok(())
+        })?;
+        clear_signal_handlers();
+        self.strategy.write().take();
+        info!("Server stopped");
+        Ok(())
     }
 }

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

@@ -0,0 +1,9 @@
+#[derive(Debug, Clone)]
+pub enum LifecycleEvent {
+    Start,
+    Shutdown,
+    Restart,
+    IncreaseWorkers,
+    DecreaseWorkers,
+    ForceShutdown,
+}