itsi-scheduler 0.1.5 → 0.1.19

data/ext/itsi_server/src/services/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/services/itsi_http_service.rs

@@ -0,0 +1,239 @@
+use crate::default_responses::{NOT_FOUND_RESPONSE, TIMEOUT_RESPONSE};
+use crate::ruby_types::itsi_server::itsi_server_config::ServerParams;
+use crate::server::binds::listener::ListenerInfo;
+use crate::server::http_message_types::{ConversionExt, HttpResponse, RequestExt, ResponseFormat};
+use crate::server::lifecycle_event::LifecycleEvent;
+use crate::server::middleware_stack::MiddlewareLayer;
+use crate::server::request_job::RequestJob;
+use crate::server::serve_strategy::single_mode::RunningPhase;
+use crate::server::signal::send_lifecycle_event;
+use chrono;
+use chrono::Local;
+use either::Either;
+use http::header::ACCEPT_ENCODING;
+use http::{HeaderValue, Request};
+use hyper::body::Incoming;
+use hyper::service::Service;
+use itsi_error::ItsiError;
+use regex::Regex;
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::OnceLock;
+
+use std::{future::Future, ops::Deref, pin::Pin, sync::Arc};
+use tokio::sync::watch::{self};
+use tokio::time::timeout;
+
+#[derive(Clone)]
+pub struct ItsiHttpService {
+    pub inner: Arc<ItsiHttpServiceInner>,
+}
+
+impl Deref for ItsiHttpService {
+    type Target = Arc<ItsiHttpServiceInner>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.inner
+    }
+}
+
+pub struct ItsiHttpServiceInner {
+    pub sender: async_channel::Sender<RequestJob>,
+    pub nonblocking_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 HttpRequestContext {
+    inner: Arc<RequestContextInner>,
+}
+
+impl Deref for HttpRequestContext {
+    type Target = Arc<RequestContextInner>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.inner
+    }
+}
+
+impl Deref for RequestContextInner {
+    type Target = ItsiHttpService;
+
+    fn deref(&self) -> &Self::Target {
+        &self.service
+    }
+}
+
+pub struct RequestContextInner {
+    pub request_id: u128,
+    pub service: ItsiHttpService,
+    pub accept: ResponseFormat,
+    pub matching_pattern: Option<Arc<Regex>>,
+    pub origin: OnceLock<Option<String>>,
+    pub response_format: OnceLock<ResponseFormat>,
+    pub start_time: chrono::DateTime<chrono::Utc>,
+    pub request: Option<Arc<Request<Incoming>>>,
+    pub request_start_time: OnceLock<chrono::DateTime<Local>>,
+    pub if_none_match: OnceLock<Option<String>>,
+    pub supported_encoding_set: Vec<HeaderValue>,
+    pub is_ruby_request: Arc<AtomicBool>,
+}
+
+impl HttpRequestContext {
+    fn new(
+        service: ItsiHttpService,
+        matching_pattern: Option<Arc<Regex>>,
+        accept: ResponseFormat,
+        supported_encoding_set: Vec<HeaderValue>,
+        is_ruby_request: Arc<AtomicBool>,
+    ) -> Self {
+        HttpRequestContext {
+            inner: Arc::new(RequestContextInner {
+                request_id: rand::random::<u128>(),
+                service,
+                matching_pattern,
+                accept,
+                origin: OnceLock::new(),
+                response_format: OnceLock::new(),
+                start_time: chrono::Utc::now(),
+                request: None,
+                request_start_time: OnceLock::new(),
+                if_none_match: OnceLock::new(),
+                supported_encoding_set,
+                is_ruby_request,
+            }),
+        }
+    }
+
+    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 short_request_id(&self) -> String {
+        format!("{:016x}", self.inner.request_id & 0xffff_ffff_ffff_ffff)
+    }
+
+    pub fn request_id(&self) -> String {
+        format!("{:016x}", self.inner.request_id)
+    }
+
+    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)
+    }
+
+    pub fn set_response_format(&self, format: ResponseFormat) {
+        self.inner.response_format.set(format).unwrap()
+    }
+
+    pub fn response_format(&self) -> &ResponseFormat {
+        self.inner.response_format.get().unwrap()
+    }
+}
+
+impl Service<Request<Incoming>> for ItsiHttpService {
+    type Response = HttpResponse;
+    type Error = ItsiError;
+    type Future = Pin<Box<dyn Future<Output = itsi_error::Result<HttpResponse>> + Send>>;
+
+    fn call(&self, req: Request<Incoming>) -> Self::Future {
+        let params = self.server_params.clone();
+        let self_clone = self.clone();
+        let mut req = req.limit();
+        let accept: ResponseFormat = req.accept().into();
+        let accept_clone = accept.clone();
+        let is_single_mode = self.server_params.workers == 1;
+        let supported_encoding_set = req
+            .headers()
+            .get_all(ACCEPT_ENCODING)
+            .into_iter()
+            .cloned()
+            .collect::<Vec<_>>();
+        let request_timeout = self.server_params.request_timeout;
+        let is_ruby_request = Arc::new(AtomicBool::new(false));
+        let irr_clone = is_ruby_request.clone();
+        let service_future = async move {
+            let mut resp: Option<HttpResponse> = None;
+            let (stack, matching_pattern) = params.middleware.get().unwrap().stack_for(&req)?;
+
+            let mut context = HttpRequestContext::new(
+                self_clone,
+                matching_pattern,
+                accept_clone.clone(),
+                supported_encoding_set,
+                irr_clone,
+            );
+            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 Ok(NOT_FOUND_RESPONSE.to_http_response(accept_clone).await),
+            };
+
+            for elm in stack.iter().rev().skip(stack.len() - depth - 1) {
+                resp = elm.after(resp, &mut context).await;
+            }
+
+            Ok(resp)
+        };
+
+        Box::pin(async move {
+            if let Some(timeout_duration) = request_timeout {
+                match timeout(timeout_duration, service_future).await {
+                    Ok(result) => result,
+                    Err(_) => {
+                        // If we're still running Ruby at this point, we can't just kill the
+                        // thread as it might be in a critical section.
+                        // Instead we must ask the worker to hot restart.
+                        if is_ruby_request.load(Ordering::Relaxed) {
+                            if is_single_mode {
+                                // If we're in single mode, re-exec the whole process
+                                send_lifecycle_event(LifecycleEvent::Restart);
+                            } else {
+                                // Otherwise we can shutdown the worker and rely on the master to restart it
+                                send_lifecycle_event(LifecycleEvent::Shutdown);
+                            }
+                        }
+                        Ok(TIMEOUT_RESPONSE.to_http_response(accept).await)
+                    }
+                }
+            } else {
+                service_future.await
+            }
+        })
+    }
+}