itsi-scheduler 0.1.5 → 0.2.2

data/ext/itsi_server/src/services/itsi_http_service.rs

@@ -0,0 +1,257 @@
+use crate::default_responses::{NOT_FOUND_RESPONSE, TIMEOUT_RESPONSE};
+use crate::ruby_types::itsi_server::itsi_server_config::{ItsiServerTokenPreference, 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 tracing::error;
+
+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()
+    }
+}
+
+const SERVER_TOKEN_VERSION: HeaderValue =
+    HeaderValue::from_static(concat!("Itsi/", env!("CARGO_PKG_VERSION")));
+const SERVER_TOKEN_NAME: HeaderValue = HeaderValue::from_static("Itsi");
+
+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) => {
+                        error!("Middleware error: {}", e);
+                        break;
+                    }
+                }
+            }
+
+            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;
+            }
+
+            match params.itsi_server_token_preference {
+                ItsiServerTokenPreference::Version => {
+                    resp.headers_mut().insert("Server", SERVER_TOKEN_VERSION);
+                }
+                ItsiServerTokenPreference::Name => {
+                    resp.headers_mut().insert("Server", SERVER_TOKEN_NAME);
+                }
+                ItsiServerTokenPreference::None => {}
+            }
+
+            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
+            }
+        })
+    }
+}