itsi-scheduler 0.1.5 → 0.2.2

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

@@ -1,59 +1,64 @@
-use std::sync::{atomic::AtomicI8, LazyLock};
+use std::sync::{
+    atomic::{AtomicBool, AtomicI8},
+    LazyLock,
+};
 
 use nix::libc::{self, sighandler_t};
 use tokio::sync::{self, broadcast};
 
 use super::lifecycle_event::LifecycleEvent;
 
+pub static SIGINT_COUNT: AtomicI8 = AtomicI8::new(0);
+pub static SHUTDOWN_REQUESTED: AtomicBool = AtomicBool::new(false);
 pub static SIGNAL_HANDLER_CHANNEL: LazyLock<(
     broadcast::Sender<LifecycleEvent>,
     broadcast::Receiver<LifecycleEvent>,
 )> = LazyLock::new(|| sync::broadcast::channel(5));
 
-pub static SIGINT_COUNT: AtomicI8 = AtomicI8::new(0);
+pub fn send_lifecycle_event(event: LifecycleEvent) {
+    SIGNAL_HANDLER_CHANNEL.0.send(event).ok();
+}
+
 fn receive_signal(signum: i32, _: sighandler_t) {
     SIGINT_COUNT.fetch_add(-1, std::sync::atomic::Ordering::SeqCst);
-    match signum {
+    let event = match signum {
         libc::SIGTERM | libc::SIGINT => {
+            SHUTDOWN_REQUESTED.store(true, std::sync::atomic::Ordering::SeqCst);
             SIGINT_COUNT.fetch_add(2, std::sync::atomic::Ordering::SeqCst);
             if SIGINT_COUNT.load(std::sync::atomic::Ordering::SeqCst) < 2 {
-                SIGNAL_HANDLER_CHANNEL.0.send(LifecycleEvent::Shutdown).ok();
+                Some(LifecycleEvent::Shutdown)
             } else {
                 // Not messing about. Force shutdown.
-                SIGNAL_HANDLER_CHANNEL
-                    .0
-                    .send(LifecycleEvent::ForceShutdown)
-                    .ok();
+                Some(LifecycleEvent::ForceShutdown)
             }
         }
-        libc::SIGUSR1 => {
-            SIGNAL_HANDLER_CHANNEL.0.send(LifecycleEvent::Restart).ok();
-        }
-        libc::SIGTTIN => {
-            SIGNAL_HANDLER_CHANNEL
-                .0
-                .send(LifecycleEvent::IncreaseWorkers)
-                .ok();
-        }
-        libc::SIGTTOU => {
-            SIGNAL_HANDLER_CHANNEL
-                .0
-                .send(LifecycleEvent::DecreaseWorkers)
-                .ok();
-        }
-        _ => {}
+        libc::SIGUSR2 => Some(LifecycleEvent::PrintInfo),
+        libc::SIGUSR1 => Some(LifecycleEvent::Restart),
+        libc::SIGHUP => Some(LifecycleEvent::Reload),
+        libc::SIGTTIN => Some(LifecycleEvent::IncreaseWorkers),
+        libc::SIGTTOU => Some(LifecycleEvent::DecreaseWorkers),
+        libc::SIGCHLD => Some(LifecycleEvent::ChildTerminated),
+        _ => None,
+    };
+
+    if let Some(event) = event {
+        send_lifecycle_event(event);
     }
 }
 
 pub fn reset_signal_handlers() -> bool {
     SIGINT_COUNT.store(0, std::sync::atomic::Ordering::SeqCst);
+    SHUTDOWN_REQUESTED.store(false, std::sync::atomic::Ordering::SeqCst);
+
     unsafe {
         libc::signal(libc::SIGTERM, receive_signal as usize);
         libc::signal(libc::SIGINT, receive_signal as usize);
-        libc::signal(libc::SIGUSR1, receive_signal as usize);
         libc::signal(libc::SIGUSR2, receive_signal as usize);
+        libc::signal(libc::SIGUSR1, receive_signal as usize);
+        libc::signal(libc::SIGHUP, receive_signal as usize);
         libc::signal(libc::SIGTTIN, receive_signal as usize);
         libc::signal(libc::SIGTTOU, receive_signal as usize);
+        libc::signal(libc::SIGCHLD, receive_signal as usize);
     }
     true
 }
@@ -62,9 +67,11 @@ pub fn clear_signal_handlers() {
     unsafe {
         libc::signal(libc::SIGTERM, libc::SIG_DFL);
         libc::signal(libc::SIGINT, libc::SIG_DFL);
-        libc::signal(libc::SIGUSR1, libc::SIG_DFL);
         libc::signal(libc::SIGUSR2, libc::SIG_DFL);
+        libc::signal(libc::SIGUSR1, libc::SIG_DFL);
+        libc::signal(libc::SIGHUP, libc::SIG_DFL);
         libc::signal(libc::SIGTTIN, libc::SIG_DFL);
         libc::signal(libc::SIGTTOU, libc::SIG_DFL);
+        libc::signal(libc::SIGCHLD, libc::SIG_DFL);
     }
 }

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

@@ -0,0 +1,107 @@
+use bytes::Buf;
+use hyper::body::Body;
+use hyper::body::Frame;
+use hyper::body::SizeHint;
+use std::error::Error;
+use std::fmt;
+use std::ops::Deref;
+use std::pin::Pin;
+use std::sync::atomic::AtomicUsize;
+use std::sync::atomic::Ordering;
+use std::task::Context;
+use std::task::Poll;
+use tracing::debug;
+
+/// Custom error to indicate that the maximum body size was exceeded.
+#[derive(Debug)]
+pub struct MaxBodySizeReached;
+impl fmt::Display for MaxBodySizeReached {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "Maximum body size reached")
+    }
+}
+
+impl Error for MaxBodySizeReached {}
+
+#[derive(Debug)]
+pub struct SizeLimitedIncoming<B> {
+    pub inner: B,
+    pub limit: AtomicUsize,
+    current: usize,
+}
+
+impl<B> Deref for SizeLimitedIncoming<B> {
+    type Target = B;
+
+    fn deref(&self) -> &Self::Target {
+        &self.inner
+    }
+}
+
+impl<B> SizeLimitedIncoming<B> {
+    pub fn new(inner: B) -> Self {
+        Self {
+            inner,
+            limit: AtomicUsize::new(usize::MAX),
+            current: 0,
+        }
+    }
+}
+
+impl<B> Body for SizeLimitedIncoming<B>
+where
+    B: Body + Unpin,
+    B::Data: Buf,
+    // Ensure that the inner error converts into our boxed error type.
+    B::Error: Into<Box<dyn Error + Send + Sync>>,
+{
+    type Data = B::Data;
+    type Error = Box<dyn Error + Send + Sync>;
+
+    fn poll_frame(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
+        // Pin the inner body.
+        let inner = Pin::new(&mut self.inner);
+        match inner.poll_frame(cx) {
+            Poll::Ready(Some(Ok(frame))) => {
+                // Use public methods since we cannot match on the private enum.
+                if frame.is_data() {
+                    match frame.into_data() {
+                        Ok(data) => {
+                            let len = data.remaining();
+                            self.current += len;
+                            debug!(
+                              target: "option::max_body",
+                              "current: {}, limit: {}",
+                              self.current, self.limit.load(Ordering::Relaxed)
+                            );
+                            if self.current > self.limit.load(Ordering::Relaxed) {
+                                Poll::Ready(Some(Err(Box::new(MaxBodySizeReached))))
+                            } else {
+                                Poll::Ready(Some(Ok(Frame::data(data))))
+                            }
+                        }
+                        // Should not occur if is_data() was true, but pass through if it does.
+                        Err(frame) => Poll::Ready(Some(Ok(frame))),
+                    }
+                } else {
+                    // For non-data frames (e.g. trailers), just pass them along.
+                    Poll::Ready(Some(Ok(frame)))
+                }
+            }
+            Poll::Ready(Some(Err(e))) => Poll::Ready(Some(Err(e.into()))),
+            Poll::Ready(None) => Poll::Ready(None),
+            Poll::Pending => Poll::Pending,
+        }
+    }
+
+    fn is_end_stream(&self) -> bool {
+        self.inner.is_end_stream()
+    }
+
+    fn size_hint(&self) -> SizeHint {
+        self.inner.size_hint()
+    }
+}