itsi-server 0.1.1 → 0.1.8

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

@@ -1,56 +1,111 @@
 use super::{
     bind::Bind,
-    listener::{Listener, SockAddr},
+    listener::Listener,
+    serve_strategy::{cluster_mode::ClusterMode, single_mode::SingleMode},
+    signal::{clear_signal_handlers, reset_signal_handlers, 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;
+use itsi_tracing::error;
 use magnus::{
+    block::Proc,
     error::Result,
     scan_args::{get_kwargs, scan_args, Args, KwArgs},
-    value::{Opaque, ReprValue},
-    RHash, Ruby, Value,
+    value::{InnerValue, Opaque, ReprValue},
+    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 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: Opaque<Value>,
     #[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>,
+}
+
+#[derive(Debug)]
+pub enum RequestJob {
+    ProcessRequest(ItsiRequest),
+    Shutdown,
 }
 
 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>,
+        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>,
         );
 
-        let scan_args: Args<(), (), (), (), RHash, ()> = scan_args(args)?;
-        let args: KwArgs<(Value,), OptionalArgs, ()> = get_kwargs(
-            scan_args.keywords,
+        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(),
             &["app"],
             &[
                 "workers",
@@ -58,125 +113,125 @@ impl Server {
                 "shutdown_timeout",
                 "script_name",
                 "binds",
+                "before_fork",
+                "after_fork",
+                "scheduler_class",
+                "stream_body",
             ],
         )?;
-        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 args2: KwArgs<(), ArgSet2, ()> = get_kwargs(
+            scan_args
+                .keywords
+                .funcall::<_, _, RHash>("slice", (Symbol::new("worker_memory_limit"),))
+                .unwrap(),
+            &[],
+            &["worker_memory_limit"],
+        )?;
+
+        let config = ServerConfig {
+            app: Opaque::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>>>()?,
             ),
+            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,
         };
-        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())
-            }
-            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)
+        if let Some(scheduler_class) = args1.optional.7 {
+            info!(scheduler_class, fiber_scheduler = true);
         } else {
-            let mut response = Response::new(BoxBody::new(Empty::new()));
-            *response.status_mut() = StatusCode::BAD_REQUEST;
-            Ok(response)
+            info!(fiber_scheduler = false);
         }
+
+        Ok(Server {
+            config: Arc::new(config),
+        })
     }
 
-    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);
-                            }
-                        });
-                    }
-                });
+    #[instrument(name = "Bind", skip_all, fields(binds=format!("{:?}", self.config.binds.lock())))]
+    pub(crate) fn 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<ServeStrategy> {
+        let server = Arc::new(self);
+
+        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(),
+            )))
+        };
+        Ok(strategy)
+    }
+
+    pub fn start(&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() {
+                error!("Error running server: {}", e);
+                strategy.stop()?;
             }
-            while let Some(_res) = set.join_next().await {}
-        })
+            drop(strategy);
+            Ok(())
+        })?;
+        clear_signal_handlers();
+        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,
+}