itsi-scheduler 0.1.11 → 0.1.12

data/ext/itsi_server/src/ruby_types/itsi_server/itsi_server_config.rs

@@ -0,0 +1,355 @@
+use super::file_watcher::{self};
+use crate::{
+    ruby_types::ITSI_SERVER_CONFIG,
+    server::{bind::Bind, listener::Listener, middleware_stack::MiddlewareSet},
+};
+use derive_more::Debug;
+use itsi_rb_helpers::{call_with_gvl, print_rb_backtrace, HeapVal, HeapValue};
+use itsi_tracing::set_level;
+use magnus::{
+    block::Proc,
+    error::Result,
+    value::{LazyId, ReprValue},
+    RArray, RHash, Ruby, Symbol, Value,
+};
+use nix::{
+    fcntl::{fcntl, FcntlArg, FdFlag},
+    unistd::{close, dup},
+};
+use parking_lot::{Mutex, RwLock};
+use std::{
+    collections::HashMap,
+    os::fd::{AsRawFd, OwnedFd, RawFd},
+    path::PathBuf,
+    sync::{Arc, OnceLock},
+};
+
+static DEFAULT_BIND: &str = "http://localhost:3000";
+static ID_BUILD_CONFIG: LazyId = LazyId::new("build_config");
+static ID_RELOAD_EXEC: LazyId = LazyId::new("reload_exec");
+
+#[derive(Debug, Clone)]
+pub struct ItsiServerConfig {
+    pub cli_params: Arc<HeapValue<RHash>>,
+    pub itsifile_path: Option<PathBuf>,
+    pub itsi_config_proc: Arc<Option<HeapValue<Proc>>>,
+    #[debug(skip)]
+    pub server_params: Arc<RwLock<Arc<ServerParams>>>,
+    pub watcher_fd: Arc<Option<OwnedFd>>,
+}
+
+#[derive(Debug)]
+pub struct ServerParams {
+    /// Cluster params
+    pub workers: u8,
+    pub worker_memory_limit: Option<u64>,
+    pub silence: bool,
+    pub shutdown_timeout: f64,
+    pub hooks: HashMap<String, HeapValue<Proc>>,
+    pub preload: bool,
+
+    pub notify_watchers: Option<Vec<(String, Vec<Vec<String>>)>>,
+    /// Worker params
+    pub threads: u8,
+    pub script_name: String,
+    pub streamable_body: bool,
+    pub multithreaded_reactor: bool,
+    pub scheduler_class: Option<String>,
+    pub oob_gc_responses_threshold: Option<u64>,
+    pub middleware_loader: HeapValue<Proc>,
+    pub default_app_loader: HeapValue<Proc>,
+    pub middleware: OnceLock<MiddlewareSet>,
+    pub binds: Vec<Bind>,
+    #[debug(skip)]
+    pub(crate) listeners: Mutex<Vec<Listener>>,
+    listener_info: Mutex<HashMap<String, i32>>,
+}
+
+impl ServerParams {
+    pub fn preload_ruby(self: &Arc<Self>) -> Result<()> {
+        call_with_gvl(|ruby| -> Result<()> {
+            if self
+                .scheduler_class
+                .as_ref()
+                .is_some_and(|t| t == "Itsi::Scheduler")
+            {
+                ruby.require("itsi/scheduler")?;
+            }
+            let default_app: HeapVal = self.default_app_loader.call::<_, Value>(())?.into();
+            let middleware = MiddlewareSet::new(
+                self.middleware_loader
+                    .call::<_, Option<Value>>(())
+                    .inspect_err(|e| {
+                        if let Some(err_value) = e.value() {
+                            print_rb_backtrace(err_value);
+                        }
+                    })?
+                    .map(|mw| mw.into()),
+                default_app,
+            )?;
+            self.middleware.set(middleware).map_err(|_| {
+                magnus::Error::new(
+                    magnus::exception::runtime_error(),
+                    "Failed to set middleware",
+                )
+            })?;
+            Ok(())
+        })?;
+        Ok(())
+    }
+
+    fn from_rb_hash(rb_param_hash: RHash) -> Result<ServerParams> {
+        let workers = rb_param_hash
+            .fetch::<_, Option<u8>>("workers")?
+            .unwrap_or(num_cpus::get() as u8);
+        let worker_memory_limit: Option<u64> = rb_param_hash.fetch("worker_memory_limit")?;
+        let silence: bool = rb_param_hash.fetch("silence")?;
+        let multithreaded_reactor: bool = rb_param_hash.fetch("multithreaded_reactor")?;
+        let shutdown_timeout: f64 = rb_param_hash.fetch("shutdown_timeout")?;
+
+        let hooks: Option<RHash> = rb_param_hash.fetch("hooks")?;
+        let hooks = hooks
+            .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 preload: bool = rb_param_hash.fetch("preload")?;
+        let notify_watchers: Option<Vec<(String, Vec<Vec<String>>)>> =
+            rb_param_hash.fetch("notify_watchers")?;
+        let threads: u8 = rb_param_hash.fetch("threads")?;
+        let script_name: String = rb_param_hash.fetch("script_name")?;
+        let streamable_body: bool = rb_param_hash.fetch("streamable_body")?;
+        let scheduler_class: Option<String> = rb_param_hash.fetch("scheduler_class")?;
+        let oob_gc_responses_threshold: Option<u64> =
+            rb_param_hash.fetch("oob_gc_responses_threshold")?;
+        let middleware_loader: Proc = rb_param_hash.fetch("middleware_loader")?;
+        let default_app_loader: Proc = rb_param_hash.fetch("default_app_loader")?;
+        let log_level: Option<String> = rb_param_hash.fetch("log_level")?;
+
+        if let Some(level) = log_level {
+            set_level(&level);
+        }
+
+        let binds: Option<Vec<String>> = rb_param_hash.fetch("binds")?;
+        let binds = binds
+            .unwrap_or_else(|| vec![DEFAULT_BIND.to_string()])
+            .into_iter()
+            .map(|s| s.parse())
+            .collect::<itsi_error::Result<Vec<Bind>>>()?;
+
+        let listeners = if let Some(preexisting_listeners) =
+            rb_param_hash.delete::<_, Option<String>>("listeners")?
+        {
+            let bind_to_fd_map: HashMap<String, i32> = serde_json::from_str(&preexisting_listeners)
+                .map_err(|e| {
+                    magnus::Error::new(
+                        magnus::exception::exception(),
+                        format!("Invalid listener info: {}", e),
+                    )
+                })?;
+
+            binds
+                .iter()
+                .cloned()
+                .map(|bind| {
+                    if let Some(fd) = bind_to_fd_map.get(&bind.listener_address_string()) {
+                        Listener::inherit_fd(bind, *fd)
+                    } else {
+                        Listener::try_from(bind)
+                    }
+                })
+                .collect::<std::result::Result<Vec<Listener>, _>>()?
+                .into_iter()
+                .collect::<Vec<_>>()
+        } else {
+            binds
+                .iter()
+                .cloned()
+                .map(Listener::try_from)
+                .collect::<std::result::Result<Vec<Listener>, _>>()?
+                .into_iter()
+                .collect::<Vec<_>>()
+        };
+
+        let listener_info = listeners
+            .iter()
+            .map(|listener| {
+                listener.handover().map_err(|e| {
+                    magnus::Error::new(magnus::exception::runtime_error(), e.to_string())
+                })
+            })
+            .collect::<Result<HashMap<String, i32>>>()?;
+
+        Ok(ServerParams {
+            workers,
+            worker_memory_limit,
+            silence,
+            multithreaded_reactor,
+            shutdown_timeout,
+            hooks,
+            preload,
+            notify_watchers,
+            threads,
+            script_name,
+            streamable_body,
+            scheduler_class,
+            oob_gc_responses_threshold,
+            binds,
+            listener_info: Mutex::new(listener_info),
+            listeners: Mutex::new(listeners),
+            middleware_loader: middleware_loader.into(),
+            default_app_loader: default_app_loader.into(),
+            middleware: OnceLock::new(),
+        })
+    }
+}
+
+impl ItsiServerConfig {
+    pub fn new(
+        ruby: &Ruby,
+        cli_params: RHash,
+        itsifile_path: Option<PathBuf>,
+        itsi_config_proc: Option<Proc>,
+    ) -> Result<Self> {
+        let itsi_config_proc = Arc::new(itsi_config_proc.map(HeapValue::from));
+        let server_params = Self::combine_params(
+            ruby,
+            cli_params,
+            itsifile_path.as_ref(),
+            itsi_config_proc.clone(),
+        )?;
+        cli_params.delete::<_, Value>(Symbol::new("listeners"))?;
+
+        let watcher_fd = if let Some(watchers) = server_params.notify_watchers.clone() {
+            file_watcher::watch_groups(watchers)?
+        } else {
+            None
+        };
+
+        Ok(ItsiServerConfig {
+            cli_params: Arc::new(cli_params.into()),
+            server_params: RwLock::new(server_params.clone()).into(),
+            itsi_config_proc,
+            itsifile_path,
+            watcher_fd: watcher_fd.into(),
+        })
+    }
+
+    /// Reload
+    pub fn reload(self: Arc<Self>, cluster_worker: bool) -> Result<bool> {
+        let server_params = call_with_gvl(|ruby| {
+            Self::combine_params(
+                &ruby,
+                self.cli_params.cloned(),
+                self.itsifile_path.as_ref(),
+                self.itsi_config_proc.clone(),
+            )
+        })?;
+
+        let is_single_mode = self.server_params.read().workers == 1;
+
+        let requires_exec = if !is_single_mode && !server_params.preload {
+            // In cluster mode children are cycled during a reload
+            // and if preload is disabled, will get a clean memory slate,
+            // so we don't need to exec.
+            false
+        } else {
+            // In non-cluster mode, or when preloading is enabled, we shouldn't try to
+            // reload inside the existing process (as new code may conflict with old),
+            // and should re-exec instead.
+            true
+        };
+
+        *self.server_params.write() = server_params.clone();
+        Ok(requires_exec && (cluster_worker || is_single_mode))
+    }
+
+    fn combine_params(
+        ruby: &Ruby,
+        cli_params: RHash,
+        itsifile_path: Option<&PathBuf>,
+        itsi_config_proc: Arc<Option<HeapValue<Proc>>>,
+    ) -> Result<Arc<ServerParams>> {
+        let inner = itsi_config_proc
+            .as_ref()
+            .clone()
+            .map(|hv| hv.clone().inner());
+        let rb_param_hash: RHash = ruby.get_inner_ref(&ITSI_SERVER_CONFIG).funcall(
+            *ID_BUILD_CONFIG,
+            (cli_params, itsifile_path.cloned(), inner),
+        )?;
+        Ok(Arc::new(ServerParams::from_rb_hash(rb_param_hash)?))
+    }
+
+    fn clear_cloexec(fd: RawFd) -> nix::Result<()> {
+        let current_flags = fcntl(fd, FcntlArg::F_GETFD)?;
+        let mut flags = FdFlag::from_bits_truncate(current_flags);
+        // Remove the FD_CLOEXEC flag
+        flags.remove(FdFlag::FD_CLOEXEC);
+        // Set the new flags back on the file descriptor
+        fcntl(fd, FcntlArg::F_SETFD(flags))?;
+        Ok(())
+    }
+
+    pub fn dup_fds(self: &Arc<Self>) -> Result<()> {
+        let binding = self.server_params.read();
+        let mut listener_info_guard = binding.listener_info.lock();
+        let dupped_fd_map = listener_info_guard
+            .iter()
+            .map(|(str, fd)| {
+                let dupped_fd = dup(*fd).map_err(|errno| {
+                    magnus::Error::new(
+                        magnus::exception::exception(),
+                        format!("Errno {} while trying to dup {}", errno, fd),
+                    )
+                })?;
+                Self::clear_cloexec(dupped_fd).map_err(|e| {
+                    magnus::Error::new(
+                        magnus::exception::exception(),
+                        format!("Failed to clear cloexec flag for fd {}: {}", dupped_fd, e),
+                    )
+                })?;
+                Ok((str.clone(), dupped_fd))
+            })
+            .collect::<Result<HashMap<String, i32>>>()?;
+        *listener_info_guard = dupped_fd_map;
+        Ok(())
+    }
+
+    pub fn stop_watcher(self: &Arc<Self>) -> Result<()> {
+        if let Some(r_fd) = self.watcher_fd.as_ref() {
+            close(r_fd.as_raw_fd()).ok();
+        }
+        Ok(())
+    }
+
+    pub fn reload_exec(self: &Arc<Self>) -> Result<()> {
+        let listener_json =
+            serde_json::to_string(&self.server_params.read().listener_info.lock().clone())
+                .map_err(|e| {
+                    magnus::Error::new(
+                        magnus::exception::exception(),
+                        format!("Invalid listener info: {}", e),
+                    )
+                })?;
+
+        self.stop_watcher()?;
+        call_with_gvl(|ruby| -> Result<()> {
+            ruby.get_inner_ref(&ITSI_SERVER_CONFIG)
+                .funcall::<_, _, Value>(*ID_RELOAD_EXEC, (listener_json,))?;
+            Ok(())
+        })?;
+        Ok(())
+    }
+}

data/ext/itsi_server/src/ruby_types/itsi_server.rs

@@ -0,0 +1,82 @@
+use crate::server::{
+    serve_strategy::{cluster_mode::ClusterMode, single_mode::SingleMode, ServeStrategy},
+    signal::{clear_signal_handlers, reset_signal_handlers, send_shutdown_event},
+};
+use itsi_rb_helpers::{call_without_gvl, print_rb_backtrace};
+use itsi_server_config::ItsiServerConfig;
+use itsi_tracing::{error, run_silently};
+use magnus::{block::Proc, error::Result, RHash, Ruby};
+use parking_lot::Mutex;
+use std::{path::PathBuf, sync::Arc};
+use tracing::{info, instrument};
+mod file_watcher;
+pub mod itsi_server_config;
+#[magnus::wrap(class = "Itsi::Server", free_immediately, size)]
+#[derive(Clone)]
+pub struct ItsiServer {
+    pub config: Arc<Mutex<Arc<ItsiServerConfig>>>,
+}
+
+impl ItsiServer {
+    pub fn new(
+        ruby: &Ruby,
+        cli_params: RHash,
+        itsifile_path: Option<PathBuf>,
+        itsi_config_proc: Option<Proc>,
+    ) -> Result<Self> {
+        Ok(Self {
+            config: Arc::new(Mutex::new(Arc::new(ItsiServerConfig::new(
+                ruby,
+                cli_params,
+                itsifile_path,
+                itsi_config_proc,
+            )?))),
+        })
+    }
+
+    pub fn stop(&self) -> Result<()> {
+        send_shutdown_event();
+        Ok(())
+    }
+
+    #[instrument(skip(self))]
+    pub fn start(&self) -> Result<()> {
+        let result = if self.config.lock().server_params.read().silence {
+            run_silently(|| self.build_and_run_strategy())
+        } else {
+            info!("Itsi - Rolling into action. 💨 ⚪ ");
+            self.build_and_run_strategy()
+        };
+        if let Err(e) = result {
+            if let Some(err_value) = e.value() {
+                print_rb_backtrace(err_value);
+            }
+            return Err(e);
+        }
+        Ok(())
+    }
+
+    pub(crate) fn build_strategy(&self) -> Result<ServeStrategy> {
+        let server_config = self.config.lock();
+        Ok(if server_config.server_params.read().workers > 1 {
+            ServeStrategy::Cluster(Arc::new(ClusterMode::new(server_config.clone())))
+        } else {
+            ServeStrategy::Single(Arc::new(SingleMode::new(server_config.clone())?))
+        })
+    }
+
+    fn build_and_run_strategy(&self) -> Result<()> {
+        reset_signal_handlers();
+        call_without_gvl(move || -> Result<()> {
+            let strategy = self.build_strategy()?;
+            if let Err(e) = strategy.clone().run() {
+                error!("Error running server: {}", e);
+                strategy.stop()?;
+            }
+            Ok(())
+        })?;
+        clear_signal_handlers();
+        info!("Server stopped");
+        Ok(())
+    }
+}

data/ext/itsi_server/src/ruby_types/mod.rs

@@ -0,0 +1,55 @@
+use magnus::{value::Lazy, Module, RClass, RModule};
+
+pub mod itsi_body_proxy;
+pub mod itsi_grpc_request;
+pub mod itsi_grpc_response;
+pub mod itsi_grpc_stream;
+pub mod itsi_http_request;
+pub mod itsi_http_response;
+pub mod itsi_server;
+
+pub static ITSI_MODULE: Lazy<RModule> = Lazy::new(|ruby| ruby.define_module("Itsi").unwrap());
+pub static ITSI_SERVER: Lazy<RClass> = Lazy::new(|ruby| {
+    ruby.get_inner(&ITSI_MODULE)
+        .define_class("Server", ruby.class_object())
+        .unwrap()
+});
+
+pub static ITSI_SERVER_CONFIG: Lazy<RModule> =
+    Lazy::new(|ruby| ruby.get_inner(&ITSI_SERVER).const_get("Config").unwrap());
+
+pub static ITSI_REQUEST: Lazy<RClass> = Lazy::new(|ruby| {
+    ruby.get_inner(&ITSI_MODULE)
+        .define_class("HttpRequest", ruby.class_object())
+        .unwrap()
+});
+
+pub static ITSI_RESPONSE: Lazy<RClass> = Lazy::new(|ruby| {
+    ruby.get_inner(&ITSI_MODULE)
+        .define_class("HttpResponse", ruby.class_object())
+        .unwrap()
+});
+
+pub static ITSI_BODY_PROXY: Lazy<RClass> = Lazy::new(|ruby| {
+    ruby.get_inner(&ITSI_MODULE)
+        .define_class("BodyProxy", ruby.class_object())
+        .unwrap()
+});
+
+pub static ITSI_GRPC_REQUEST: Lazy<RClass> = Lazy::new(|ruby| {
+    ruby.get_inner(&ITSI_MODULE)
+        .define_class("GrpcRequest", ruby.class_object())
+        .unwrap()
+});
+
+pub static ITSI_GRPC_STREAM: Lazy<RClass> = Lazy::new(|ruby| {
+    ruby.get_inner(&ITSI_MODULE)
+        .define_class("GrpcStream", ruby.class_object())
+        .unwrap()
+});
+
+pub static ITSI_GRPC_RESPONSE: Lazy<RClass> = Lazy::new(|ruby| {
+    ruby.get_inner(&ITSI_MODULE)
+        .define_class("GrpcResponse", ruby.class_object())
+        .unwrap()
+});

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

@@ -31,13 +31,21 @@ pub struct Bind {
     pub tls_config: Option<ItsiTlsAcceptor>,
 }
 
+impl Bind {
+    pub fn listener_address_string(&self) -> String {
+        match &self.address {
+            BindAddress::Ip(ip) => format!("tcp://{}:{}", ip.to_canonical(), self.port.unwrap()),
+            BindAddress::UnixSocket(path) => {
+                format!("unix://{}", path.as_path().to_str().unwrap())
+            }
+        }
+    }
+}
+
 impl std::fmt::Debug for Bind {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match &self.address {
             BindAddress::Ip(ip) => match self.protocol {
-                BindProtocol::Unix | BindProtocol::Unixs => {
-                    write!(f, "{}://{}", self.protocol, ip)
-                }
                 BindProtocol::Https if self.port == Some(443) => {
                     write!(f, "{}://{}", self.protocol, ip)
                 }
@@ -158,8 +166,8 @@ fn resolve_hostname(hostname: &str) -> Option<IpAddr> {
         .to_socket_addrs()
         .ok()?
         .find_map(|addr| {
-            if addr.is_ipv6() {
-                Some(addr.ip()) // Prefer IPv6
+            if addr.is_ipv4() {
+                Some(addr.ip()) // Prefer IPv4
             } else {
                 None
             }

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

@@ -0,0 +1,32 @@
+use std::ops::Deref;
+
+use bytes::Bytes;
+
+#[derive(Debug)]
+pub enum ByteFrame {
+    Data(Bytes),
+    End(Bytes),
+    Empty,
+}
+
+impl Deref for ByteFrame {
+    type Target = Bytes;
+
+    fn deref(&self) -> &Self::Target {
+        match self {
+            ByteFrame::Data(data) => data,
+            ByteFrame::End(data) => data,
+            ByteFrame::Empty => unreachable!(),
+        }
+    }
+}
+
+impl From<ByteFrame> for Bytes {
+    fn from(frame: ByteFrame) -> Self {
+        match frame {
+            ByteFrame::Data(data) => data,
+            ByteFrame::End(data) => data,
+            ByteFrame::Empty => unreachable!(),
+        }
+    }
+}

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

@@ -0,0 +1,74 @@
+use async_trait::async_trait;
+use redis::aio::ConnectionManager;
+use redis::{Client, RedisError, Script};
+use std::sync::Arc;
+use std::time::Duration;
+
+#[derive(Debug)]
+pub enum CacheError {
+    RedisError(RedisError),
+    // Other error variants as needed.
+}
+/// A general-purpose cache trait with an atomic “increment with timeout” operation.
+#[async_trait]
+pub trait CacheStore: Send + Sync + std::fmt::Debug {
+    /// Increments the counter associated with `key` and sets (or extends) its expiration.
+    /// Returns the new counter value.
+    async fn increment(&self, key: &str, timeout: Duration) -> Result<u64, CacheError>;
+}
+
+/// A Redis-backed cache store using an async connection manager.
+/// This uses a TLS-enabled connection when the URL is prefixed with "rediss://".
+#[derive(Clone)]
+pub struct RedisCacheStore {
+    connection: Arc<ConnectionManager>,
+}
+
+impl std::fmt::Debug for RedisCacheStore {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("RedisCacheStore").finish()
+    }
+}
+
+impl RedisCacheStore {
+    /// Constructs a new RedisCacheStore.
+    ///
+    /// Use a connection URL like "rediss://host:port" to enable TLS (with rustls under the hood).
+    /// This constructor is async because it sets up the connection manager.
+    pub async fn new(connection_url: &str) -> Result<Self, CacheError> {
+        let client = Client::open(connection_url).map_err(CacheError::RedisError)?;
+        let connection_manager = ConnectionManager::new(client)
+            .await
+            .map_err(CacheError::RedisError)?;
+        Ok(Self {
+            connection: Arc::new(connection_manager),
+        })
+    }
+}
+
+#[async_trait]
+impl CacheStore for RedisCacheStore {
+    async fn increment(&self, key: &str, timeout: Duration) -> Result<u64, CacheError> {
+        let timeout_secs = timeout.as_secs();
+        // Lua script to:
+        // 1. INCR the key.
+        // 2. If the key doesn't have a TTL, set it.
+        let script = r#"
+            local current = redis.call('INCR', KEYS[1])
+            if redis.call('TTL', KEYS[1]) < 0 then
+                redis.call('EXPIRE', KEYS[1], ARGV[1])
+            end
+            return current
+        "#;
+        let script = Script::new(script);
+        // The ConnectionManager is cloneable and can be used concurrently.
+        let mut connection = (*self.connection).clone();
+        let value: i64 = script
+            .key(key)
+            .arg(timeout_secs)
+            .invoke_async(&mut connection)
+            .await
+            .map_err(CacheError::RedisError)?;
+        Ok(value as u64)
+    }
+}