itsi-scheduler 0.1.5 → 0.1.14

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)
+    }
+}

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

@@ -0,0 +1,172 @@
+use super::listener::ListenerInfo;
+use super::middleware_stack::CompressionAlgorithm;
+use super::middleware_stack::MiddlewareLayer;
+use super::request_job::RequestJob;
+use super::serve_strategy::single_mode::RunningPhase;
+use super::types::HttpRequest;
+use super::types::HttpResponse;
+use crate::ruby_types::itsi_server::itsi_server_config::ServerParams;
+use chrono;
+use chrono::Local;
+use either::Either;
+use hyper::service::Service;
+use itsi_error::ItsiError;
+use regex::Regex;
+use std::sync::OnceLock;
+use std::{future::Future, ops::Deref, pin::Pin, sync::Arc};
+use tokio::sync::watch::{self};
+
+#[derive(Clone)]
+pub struct ItsiService {
+    pub inner: Arc<IstiServiceInner>,
+}
+
+impl Deref for ItsiService {
+    type Target = Arc<IstiServiceInner>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.inner
+    }
+}
+
+pub struct IstiServiceInner {
+    pub sender: async_channel::Sender<RequestJob>,
+    pub server_params: Arc<ServerParams>,
+    pub listener: Arc<ListenerInfo>,
+    pub addr: String,
+    pub shutdown_channel: watch::Receiver<RunningPhase>,
+}
+
+#[derive(Clone)]
+pub struct RequestContext {
+    inner: Arc<RequestContextInner>,
+}
+
+impl Deref for RequestContext {
+    type Target = Arc<RequestContextInner>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.inner
+    }
+}
+
+impl Deref for RequestContextInner {
+    type Target = ItsiService;
+
+    fn deref(&self) -> &Self::Target {
+        &self.service
+    }
+}
+
+pub struct RequestContextInner {
+    pub request_id: i128,
+    pub service: ItsiService,
+    pub matching_pattern: Option<Arc<Regex>>,
+    pub compression_method: OnceLock<CompressionAlgorithm>,
+    pub origin: OnceLock<Option<String>>,
+    pub start_time: chrono::DateTime<chrono::Utc>,
+    pub request: Option<Arc<HttpRequest>>,
+    pub request_start_time: OnceLock<chrono::DateTime<Local>>,
+    pub if_none_match: OnceLock<Option<String>>,
+    pub etag_value: OnceLock<Option<String>>,
+}
+
+impl RequestContext {
+    fn new(service: ItsiService, matching_pattern: Option<Arc<Regex>>) -> Self {
+        RequestContext {
+            inner: Arc::new(RequestContextInner {
+                request_id: rand::random::<i128>(),
+                service,
+                matching_pattern,
+                compression_method: OnceLock::new(),
+                origin: OnceLock::new(),
+                start_time: chrono::Utc::now(),
+                request: None,
+                request_start_time: OnceLock::new(),
+                if_none_match: OnceLock::new(),
+                etag_value: OnceLock::new(),
+            }),
+        }
+    }
+
+    pub fn set_compression_method(&self, method: CompressionAlgorithm) {
+        self.inner.compression_method.set(method).unwrap();
+    }
+
+    pub fn set_origin(&self, origin: Option<String>) {
+        self.inner.origin.set(origin).unwrap();
+    }
+
+    pub fn set_if_none_match(&self, value: Option<String>) {
+        self.inner.if_none_match.set(value).unwrap();
+    }
+
+    pub fn get_if_none_match(&self) -> Option<String> {
+        self.inner.if_none_match.get().cloned().flatten()
+    }
+
+    pub fn request_id(&self) -> String {
+        self.inner.request_id.to_string()
+    }
+
+    pub fn track_start_time(&self) {
+        self.inner
+            .request_start_time
+            .get_or_init(chrono::Local::now);
+    }
+
+    pub fn start_time(&self) -> Option<chrono::DateTime<Local>> {
+        self.inner.request_start_time.get().cloned()
+    }
+
+    pub fn get_response_time(&self) -> Option<chrono::TimeDelta> {
+        self.inner
+            .request_start_time
+            .get()
+            .map(|instant| Local::now() - instant)
+    }
+}
+
+impl Service<HttpRequest> for ItsiService {
+    type Response = HttpResponse;
+    type Error = ItsiError;
+    type Future = Pin<Box<dyn Future<Output = itsi_error::Result<HttpResponse>> + Send>>;
+
+    // This is called once per incoming Request.
+    fn call(&self, req: HttpRequest) -> Self::Future {
+        let params = self.server_params.clone();
+        let self_clone = self.clone();
+        Box::pin(async move {
+            let mut req = req;
+            let mut resp: Option<HttpResponse> = None;
+            let (stack, matching_pattern) = params.middleware.get().unwrap().stack_for(&req);
+            let mut context = RequestContext::new(self_clone, matching_pattern);
+            let mut depth = 0;
+            for (index, elm) in stack.iter().enumerate() {
+                match elm.before(req, &mut context).await {
+                    Ok(Either::Left(r)) => req = r,
+                    Ok(Either::Right(r)) => {
+                        resp = Some(r);
+                        depth = index;
+                        break;
+                    }
+                    Err(e) => return Err(e.into()),
+                }
+            }
+
+            let mut resp = match resp {
+                Some(r) => r,
+                None => {
+                    return Err(ItsiError::InternalServerError(
+                        "No response returned from middleware stack".to_string(),
+                    ))
+                }
+            };
+            for elm in stack.iter().rev().skip(stack.len() - depth - 1) {
+                resp = elm.after(resp, &mut context).await;
+            }
+
+            Ok(resp)
+        })
+    }
+}

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

@@ -3,7 +3,10 @@ pub enum LifecycleEvent {
     Start,
     Shutdown,
     Restart,
+    Reload,
     IncreaseWorkers,
     DecreaseWorkers,
     ForceShutdown,
+    PrintInfo,
+    ChildTerminated,
 }

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

@@ -1,16 +1,20 @@
 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;
 use socket2::{Domain, Protocol, Socket, Type};
+use std::fmt::Display;
 use std::net::{IpAddr, SocketAddr, TcpListener};
+use std::os::fd::{AsRawFd, FromRawFd, RawFd};
 use std::sync::Arc;
 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 +27,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 +108,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 +220,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 {
@@ -225,51 +243,103 @@ impl std::fmt::Display for SockAddr {
         }
     }
 }
+impl Display for Listener {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Listener::Tcp(listener) | Listener::TcpTls((listener, _)) => write!(
+                f,
+                "{}",
+                listener
+                    .local_addr()
+                    .map(|addr| addr.to_string())
+                    .unwrap_or_else(|_| "".to_string())
+            ),
+
+            Listener::Unix(listener) | Listener::UnixTls((listener, _)) => write!(
+                f,
+                "{}",
+                listener
+                    .local_addr()
+                    .map(|addr| addr
+                        .as_pathname()
+                        .map(|path| path.to_str().unwrap_or("").to_owned())
+                        .unwrap_or_default())
+                    .unwrap_or_else(|_| "".to_string())
+            ),
+        }
+    }
+}
 
 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),
-        };
+            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(),
+            ),
+        }
     }
-    pub fn to_tokio_listener(&self) -> TokioListener {
+
+    /// Handover information when using exec to hand over the listener to a replacement process.
+    pub fn handover(&self) -> Result<(String, i32)> {
         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::Tcp(listener) => {
+                let addr = listener.local_addr()?;
+                Ok((
+                    format!("tcp://{}:{}", addr.ip().to_canonical(), addr.port()),
+                    listener.as_raw_fd(),
+                ))
+            }
+            Listener::TcpTls((listener, _)) => {
+                let addr = listener.local_addr()?;
+                Ok((
+                    format!("tcp://{}:{}", addr.ip().to_canonical(), addr.port()),
+                    listener.as_raw_fd(),
+                ))
+            }
+            Listener::Unix(listener) => {
+                let addr = listener.local_addr()?;
+                Ok((
+                    format!("unix://{}", addr.as_pathname().unwrap().to_str().unwrap()),
+                    listener.as_raw_fd(),
+                ))
+            }
+            Listener::UnixTls((listener, _)) => {
+                let addr = listener.local_addr()?;
+                Ok((
+                    format!("unix://{}", addr.as_pathname().unwrap().to_str().unwrap()),
+                    listener.as_raw_fd(),
+                ))
+            }
+        }
+    }
+
+    pub fn inherit_fd(bind: Bind, fd: RawFd) -> Result<Self> {
+        let bound = match bind.address {
+            BindAddress::Ip(_) => match bind.protocol {
+                BindProtocol::Http => Listener::Tcp(revive_tcp_socket(fd)?),
+                BindProtocol::Https => {
+                    let tcp_listener = revive_tcp_socket(fd)?;
+                    Listener::TcpTls((tcp_listener, bind.tls_config.unwrap()))
+                }
+                _ => unreachable!(),
             },
-            Listener::UnixTls((listener, acceptor)) => TokioListener::UnixTls {
-                listener: TokioUnixListener::from_std(UnixListener::try_clone(listener).unwrap())
-                    .unwrap(),
-                acceptor: acceptor.clone(),
+            BindAddress::UnixSocket(_) => match bind.tls_config {
+                Some(tls_config) => Listener::UnixTls((revive_unix_socket(fd)?, tls_config)),
+                None => Listener::Unix(revive_unix_socket(fd)?),
             },
-        }
+        };
+        Ok(bound)
     }
 }
 
@@ -295,6 +365,27 @@ impl TryFrom<Bind> for Listener {
     }
 }
 
+fn revive_tcp_socket(fd: RawFd) -> Result<TcpListener> {
+    let socket = unsafe { Socket::from_raw_fd(fd) };
+    socket.set_reuse_port(true).ok();
+    socket.set_reuse_address(true).ok();
+    socket.set_nonblocking(true).ok();
+    socket.set_nodelay(true).ok();
+    socket.set_recv_buffer_size(262_144).ok();
+    socket.set_cloexec(true)?;
+    socket.listen(1024)?;
+    Ok(socket.into())
+}
+
+fn revive_unix_socket(fd: RawFd) -> Result<UnixListener> {
+    let socket = unsafe { Socket::from_raw_fd(fd) };
+    socket.set_nonblocking(true).ok();
+    socket.listen(1024)?;
+    socket.set_cloexec(true)?;
+
+    Ok(socket.into())
+}
+
 fn connect_tcp_socket(addr: IpAddr, port: u16) -> Result<TcpListener> {
     let domain = match addr {
         IpAddr::V4(_) => Domain::IPV4,
@@ -307,6 +398,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();
+    socket.set_only_v6(false).ok();
     socket.bind(&socket_address.into())?;
     socket.listen(1024)?;
     Ok(socket.into())
@@ -319,7 +411,6 @@ fn connect_unix_socket(path: &PathBuf) -> Result<UnixListener> {
 
     let socket_address = socket2::SockAddr::unix(path)?;
 
-    info!("Binding to {:?}", path);
     socket.bind(&socket_address)?;
     socket.listen(1024)?;