itsi-server 0.1.1 → 0.1.2

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

@@ -1,22 +1,17 @@
 use super::bind::{Bind, BindAddress};
-use super::transfer_protocol::TransferProtocol;
-use hyper_util::rt::TokioIo;
+use super::bind_protocol::BindProtocol;
+use super::io_stream::IoStream;
 use itsi_error::Result;
 use itsi_tracing::info;
 use socket2::{Domain, Protocol, Socket, Type};
-use std::net::{IpAddr, SocketAddr, TcpListener as StdTcpListener};
-use std::pin::Pin;
+use std::net::{IpAddr, SocketAddr, TcpListener};
 use std::sync::Arc;
-use std::{os::unix::net::UnixListener as StdUnixListener, path::PathBuf};
-use tokio::net::{unix, TcpListener, TcpStream, UnixListener, UnixStream};
+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_rustls::TlsAcceptor;
 
-pub(crate) trait IoStream:
-    tokio::io::AsyncRead + tokio::io::AsyncWrite + Send + Unpin
-{
-}
-impl<T: tokio::io::AsyncRead + tokio::io::AsyncWrite + Send + Unpin> IoStream for T {}
-
 pub(crate) enum Listener {
     Tcp(TcpListener),
     TcpTls((TcpListener, TlsAcceptor)),
@@ -24,149 +19,200 @@ pub(crate) enum Listener {
     UnixTls((UnixListener, TlsAcceptor)),
 }
 
-enum Stream {
-    TcpStream((TcpStream, SocketAddr)),
-    UnixStream((UnixStream, unix::SocketAddr)),
+pub(crate) enum TokioListener {
+    Tcp {
+        listener: TokioTcpListener,
+        host: String,
+        port: u16,
+    },
+    TcpTls {
+        listener: TokioTcpListener,
+        acceptor: TlsAcceptor,
+        host: String,
+        port: u16,
+    },
+    Unix {
+        listener: TokioUnixListener,
+    },
+    UnixTls {
+        listener: TokioUnixListener,
+        acceptor: TlsAcceptor,
+    },
 }
 
-#[derive(Clone, Debug)]
-pub enum SockAddr {
-    Tcp(Arc<SocketAddr>),
-    Unix(Arc<unix::SocketAddr>),
-}
-impl std::fmt::Display for SockAddr {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+impl TokioListener {
+    pub(crate) async fn accept(&self) -> Result<IoStream> {
         match self {
-            SockAddr::Tcp(socket_addr) => write!(f, "{}", socket_addr.ip().to_canonical()),
-            SockAddr::Unix(socket_addr) => match socket_addr.as_pathname() {
-                Some(path) => write!(f, "{:?}", path),
-                None => write!(f, ""),
-            },
+            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,
         }
     }
-}
 
-impl Listener {
-    pub(crate) async fn accept(&self) -> Result<(TokioIo<Pin<Box<dyn IoStream>>>, SockAddr)> {
-        match self {
-            Listener::Tcp(listener) => Listener::accept_tcp(listener).await,
-            Listener::TcpTls((listener, acceptor)) => {
-                Listener::accept_tls(listener, acceptor).await
-            }
-            Listener::Unix(stream) => Listener::accept_unix(stream).await,
-            Listener::UnixTls((listener, acceptor)) => {
-                Listener::accept_unix_tls(listener, acceptor).await
-            }
-        }
+    async fn accept_tcp(listener: &TokioTcpListener) -> Result<IoStream> {
+        let tcp_stream = listener.accept().await?;
+        Self::to_tokio_io(Stream::TcpStream(tcp_stream), None).await
+    }
+
+    async fn accept_tls(listener: &TokioTcpListener, acceptor: &TlsAcceptor) -> Result<IoStream> {
+        let tcp_stream = listener.accept().await?;
+        Self::to_tokio_io(Stream::TcpStream(tcp_stream), Some(acceptor)).await
+    }
+
+    async fn accept_unix(listener: &TokioUnixListener) -> Result<IoStream> {
+        let unix_stream = listener.accept().await?;
+        Self::to_tokio_io(Stream::UnixStream(unix_stream), None).await
+    }
+
+    async fn accept_unix_tls(
+        listener: &TokioUnixListener,
+        acceptor: &TlsAcceptor,
+    ) -> Result<IoStream> {
+        let unix_stream = listener.accept().await?;
+        Self::to_tokio_io(Stream::UnixStream(unix_stream), Some(acceptor)).await
     }
 
     async fn to_tokio_io(
         input_stream: Stream,
         tls_acceptor: Option<&TlsAcceptor>,
-    ) -> Result<(TokioIo<Pin<Box<dyn IoStream>>>, SockAddr)> {
+    ) -> Result<IoStream> {
         match tls_acceptor {
             Some(acceptor) => match input_stream {
                 Stream::TcpStream((tcp_stream, socket_address)) => {
                     match acceptor.accept(tcp_stream).await {
-                        Ok(tls_stream) => Ok((
-                            TokioIo::new(Box::pin(tls_stream) as Pin<Box<dyn IoStream>>),
-                            SockAddr::Tcp(Arc::new(socket_address)),
-                        )),
+                        Ok(tls_stream) => Ok(IoStream::TcpTls {
+                            stream: tls_stream,
+                            addr: SockAddr::Tcp(Arc::new(socket_address)),
+                        }),
                         Err(err) => Err(err.into()),
                     }
                 }
                 Stream::UnixStream((unix_stream, socket_address)) => {
                     match acceptor.accept(unix_stream).await {
-                        Ok(tls_stream) => Ok((
-                            TokioIo::new(Box::pin(tls_stream) as Pin<Box<dyn IoStream>>),
-                            SockAddr::Unix(Arc::new(socket_address)),
-                        )),
+                        Ok(tls_stream) => Ok(IoStream::UnixTls {
+                            stream: tls_stream,
+                            addr: SockAddr::Unix(Arc::new(socket_address)),
+                        }),
                         Err(err) => Err(err.into()),
                     }
                 }
             },
             None => match input_stream {
-                Stream::TcpStream((tcp_stream, socket_address)) => Ok((
-                    TokioIo::new(Box::pin(tcp_stream) as Pin<Box<dyn IoStream>>),
-                    SockAddr::Tcp(Arc::new(socket_address)),
-                )),
-                Stream::UnixStream((unix_stream, socket_address)) => Ok((
-                    TokioIo::new(Box::pin(unix_stream) as Pin<Box<dyn IoStream>>),
-                    SockAddr::Unix(Arc::new(socket_address)),
-                )),
+                Stream::TcpStream((tcp_stream, socket_address)) => Ok(IoStream::Tcp {
+                    stream: tcp_stream,
+                    addr: SockAddr::Tcp(Arc::new(socket_address)),
+                }),
+                Stream::UnixStream((unix_stream, socket_address)) => Ok(IoStream::Unix {
+                    stream: unix_stream,
+                    addr: SockAddr::Unix(Arc::new(socket_address)),
+                }),
             },
         }
     }
 
-    async fn accept_tcp(
-        listener: &TcpListener,
-    ) -> Result<(TokioIo<Pin<Box<dyn IoStream>>>, SockAddr)> {
-        let tcp_stream = listener.accept().await?;
-        Self::to_tokio_io(Stream::TcpStream(tcp_stream), None).await
-    }
-
-    async fn accept_tls(
-        listener: &TcpListener,
-        acceptor: &TlsAcceptor,
-    ) -> Result<(TokioIo<Pin<Box<dyn IoStream>>>, SockAddr)> {
-        let tcp_stream = listener.accept().await?;
-        Self::to_tokio_io(Stream::TcpStream(tcp_stream), Some(acceptor)).await
-    }
-
-    async fn accept_unix(
-        listener: &UnixListener,
-    ) -> Result<(TokioIo<Pin<Box<dyn IoStream>>>, SockAddr)> {
-        let unix_stream = listener.accept().await?;
-        Self::to_tokio_io(Stream::UnixStream(unix_stream), None).await
+    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(),
+        }
     }
 
-    async fn accept_unix_tls(
-        listener: &UnixListener,
-        acceptor: &TlsAcceptor,
-    ) -> Result<(TokioIo<Pin<Box<dyn IoStream>>>, SockAddr)> {
-        let unix_stream = listener.accept().await?;
-        Self::to_tokio_io(Stream::UnixStream(unix_stream), Some(acceptor)).await
+    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 scheme(&self) -> String {
+    pub(crate) fn host(&self) -> String {
         match self {
-            Listener::Tcp(_) => "http".to_string(),
-            Listener::TcpTls(_) => "https".to_string(),
-            Listener::Unix(_) => "http".to_string(),
-            Listener::UnixTls(_) => "https".to_string(),
+            TokioListener::Tcp { host, .. } => host.to_string(),
+            TokioListener::TcpTls { host, .. } => host.to_string(),
+            TokioListener::Unix { .. } => "unix".to_string(),
+            TokioListener::UnixTls { .. } => "unix".to_string(),
         }
     }
+}
 
-    pub(crate) fn port(&self) -> u16 {
+enum Stream {
+    TcpStream((TcpStream, SocketAddr)),
+    UnixStream((UnixStream, unix::SocketAddr)),
+}
+
+#[derive(Clone, Debug)]
+pub enum SockAddr {
+    Tcp(Arc<SocketAddr>),
+    Unix(Arc<unix::SocketAddr>),
+}
+impl std::fmt::Display for SockAddr {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
-            Listener::Tcp(listener) => listener.local_addr().unwrap().port(),
-            Listener::TcpTls((listener, _)) => listener.local_addr().unwrap().port(),
-            Listener::Unix(_) => 0,
-            Listener::UnixTls(_) => 0,
+            SockAddr::Tcp(socket_addr) => write!(f, "{}", socket_addr.ip().to_canonical()),
+            SockAddr::Unix(socket_addr) => match socket_addr.as_pathname() {
+                Some(path) => write!(f, "{:?}", path),
+                None => write!(f, ""),
+            },
         }
     }
+}
 
-    pub(crate) fn host(&self) -> String {
+impl Listener {
+    pub fn to_tokio_listener(&self) -> TokioListener {
         match self {
-            Listener::Tcp(listener) => listener.local_addr().unwrap().ip().to_string(),
-            Listener::TcpTls((listener, _)) => listener.local_addr().unwrap().ip().to_string(),
-            Listener::Unix(_) => "unix".to_string(),
-            Listener::UnixTls(_) => "unix".to_string(),
+            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::UnixTls((listener, acceptor)) => TokioListener::UnixTls {
+                listener: TokioUnixListener::from_std(UnixListener::try_clone(listener).unwrap())
+                    .unwrap(),
+                acceptor: acceptor.clone(),
+            },
         }
     }
 }
 
-impl From<Bind> for Listener {
-    fn from(bind: Bind) -> Self {
-        match bind.address {
+impl TryFrom<Bind> for Listener {
+    type Error = itsi_error::ItsiError;
+
+    fn try_from(bind: Bind) -> std::result::Result<Self, Self::Error> {
+        let bound = match bind.address {
             BindAddress::Ip(addr) => match bind.protocol {
-                TransferProtocol::Http => Listener::Tcp(
-                    TcpListener::from_std(connect_tcp_socket(addr, bind.port.unwrap())).unwrap(),
-                ),
-                TransferProtocol::Https => {
-                    let tcp_listener =
-                        TcpListener::from_std(connect_tcp_socket(addr, bind.port.unwrap()))
-                            .unwrap();
+                BindProtocol::Http => Listener::Tcp(connect_tcp_socket(addr, bind.port.unwrap())?),
+                BindProtocol::Https => {
+                    let tcp_listener = connect_tcp_socket(addr, bind.port.unwrap())?;
                     let tls_acceptor = TlsAcceptor::from(Arc::new(bind.tls_config.unwrap()));
                     Listener::TcpTls((tcp_listener, tls_acceptor))
                 }
@@ -175,44 +221,39 @@ impl From<Bind> for Listener {
             BindAddress::UnixSocket(path) => match bind.tls_config {
                 Some(tls_config) => {
                     let tls_acceptor = TlsAcceptor::from(Arc::new(tls_config));
-                    Listener::UnixTls((
-                        UnixListener::from_std(connect_unix_socket(&path)).unwrap(),
-                        tls_acceptor,
-                    ))
+                    Listener::UnixTls((connect_unix_socket(&path)?, tls_acceptor))
                 }
-                None => Listener::Unix(UnixListener::from_std(connect_unix_socket(&path)).unwrap()),
+                None => Listener::Unix(connect_unix_socket(&path)?),
             },
-        }
+        };
+        Ok(bound)
     }
 }
 
-fn connect_tcp_socket(addr: IpAddr, port: u16) -> StdTcpListener {
+fn connect_tcp_socket(addr: IpAddr, port: u16) -> Result<TcpListener> {
     let domain = match addr {
         IpAddr::V4(_) => Domain::IPV4,
         IpAddr::V6(_) => Domain::IPV6,
     };
-    let socket = Socket::new(domain, Type::STREAM, Some(Protocol::TCP)).unwrap();
+    let socket = Socket::new(domain, Type::STREAM, Some(Protocol::TCP))?;
     let socket_address: SocketAddr = SocketAddr::new(addr, port);
-    socket.set_reuse_address(true).ok();
-    socket.set_reuse_port(true).ok();
     socket.set_nonblocking(true).ok();
     socket.set_nodelay(true).ok();
     socket.set_recv_buffer_size(1_048_576).ok();
-    info!("Binding to {}", socket_address);
-    socket.bind(&socket_address.into()).unwrap();
-    socket.listen(1024).unwrap();
-    socket.into()
+    socket.bind(&socket_address.into())?;
+    socket.listen(1024)?;
+    Ok(socket.into())
 }
 
-fn connect_unix_socket(path: &PathBuf) -> StdUnixListener {
+fn connect_unix_socket(path: &PathBuf) -> Result<UnixListener> {
     let _ = std::fs::remove_file(path);
-    let socket = Socket::new(Domain::UNIX, Type::STREAM, None).unwrap();
+    let socket = Socket::new(Domain::UNIX, Type::STREAM, None)?;
     socket.set_nonblocking(true).ok();
-    let socket_address = socket2::SockAddr::unix(path).unwrap();
+    let socket_address = socket2::SockAddr::unix(path)?;
 
     info!("Binding to {:?}", path);
-    socket.bind(&socket_address).unwrap();
-    socket.listen(1024).unwrap();
+    socket.bind(&socket_address)?;
+    socket.listen(1024)?;
 
-    socket.into()
+    Ok(socket.into())
 }

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

@@ -1,5 +1,11 @@
 pub mod bind;
+pub mod bind_protocol;
+pub mod io_stream;
 pub mod itsi_server;
+pub mod lifecycle_event;
 pub mod listener;
+pub mod process_worker;
+pub mod serve_strategy;
+pub mod signal;
+pub mod thread_worker;
 pub mod tls;
-pub mod transfer_protocol;

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

@@ -0,0 +1,196 @@
+use super::serve_strategy::{cluster_mode::ClusterMode, single_mode::SingleMode};
+use itsi_error::{ItsiError, Result};
+use itsi_rb_helpers::{call_with_gvl, call_without_gvl, create_ruby_thread, fork};
+use itsi_tracing::error;
+use nix::{
+    errno::Errno,
+    sys::{
+        signal::{
+            kill,
+            Signal::{SIGKILL, SIGTERM},
+        },
+        wait::{waitpid, WaitPidFlag, WaitStatus},
+    },
+    unistd::{setpgid, Pid},
+};
+use parking_lot::Mutex;
+use std::{
+    process::{self, exit},
+    sync::Arc,
+    time::{Duration, Instant},
+};
+use sysinfo::System;
+
+use tokio::{sync::watch, time::sleep};
+use tracing::{info, instrument, warn};
+
+#[derive(Clone, Debug)]
+pub struct ProcessWorker {
+    pub worker_id: usize,
+    pub child_pid: Arc<Mutex<Option<Pid>>>,
+    pub started_at: Instant,
+}
+
+impl Default for ProcessWorker {
+    fn default() -> Self {
+        Self {
+            worker_id: 0,
+            child_pid: Arc::new(Mutex::new(None)),
+            started_at: Instant::now(),
+        }
+    }
+}
+
+impl ProcessWorker {
+    #[instrument(skip(self, cluster_template), fields(self.worker_id = %self.worker_id))]
+    pub(crate) fn boot(&self, cluster_template: Arc<ClusterMode>) -> Result<()> {
+        let child_pid = *self.child_pid.lock();
+        if let Some(pid) = child_pid {
+            if self.is_alive() {
+                if let Err(e) = kill(pid, SIGTERM) {
+                    info!("Failed to send SIGTERM to process {}: {}", pid, e);
+                }
+            }
+            *self.child_pid.lock() = None;
+        }
+
+        match call_with_gvl(|_ruby| fork(cluster_template.server.after_fork.lock().clone())) {
+            Some(pid) => {
+                *self.child_pid.lock() = Some(Pid::from_raw(pid));
+            }
+            None => {
+                if let Err(e) = setpgid(
+                    Pid::from_raw(process::id() as i32),
+                    Pid::from_raw(process::id() as i32),
+                ) {
+                    error!("Failed to set process group ID: {}", e);
+                }
+                match SingleMode::new(
+                    cluster_template.server.clone(),
+                    cluster_template.listeners.clone(),
+                    cluster_template.lifecycle_channel.clone(),
+                ) {
+                    Ok(single_mode) => {
+                        Arc::new(single_mode).run().ok();
+                    }
+                    Err(e) => {
+                        error!("Failed to boot into worker mode: {}", e);
+                    }
+                }
+                exit(0)
+            }
+        }
+        Ok(())
+    }
+
+    pub(crate) fn memory_usage(&self) -> Option<u64> {
+        if let Some(pid) = *self.child_pid.lock() {
+            let s = System::new_all();
+            if let Some(process) = s.process(sysinfo::Pid::from(pid.as_raw() as usize)) {
+                return Some(process.memory());
+            }
+        }
+        None
+    }
+
+    pub(crate) async fn reboot(&self, cluster_template: Arc<ClusterMode>) -> Result<bool> {
+        self.graceful_shutdown(cluster_template.clone()).await;
+        let self_clone = self.clone();
+        let (booted_sender, mut booted_receiver) = watch::channel(false);
+        create_ruby_thread(move || {
+            call_without_gvl(move || {
+                if self_clone.boot(cluster_template).is_ok() {
+                    booted_sender.send(true).ok()
+                } else {
+                    booted_sender.send(false).ok()
+                };
+            })
+        });
+
+        booted_receiver
+            .changed()
+            .await
+            .map_err(|_| ItsiError::InternalServerError("Failed to boot worker".to_owned()))?;
+
+        let guard = booted_receiver.borrow();
+        let result = guard.to_owned();
+        // Not very robust, we should check to see if the worker is actually listening before considering this successful.
+        sleep(Duration::from_secs(1)).await;
+        Ok(result)
+    }
+
+    pub(crate) async fn graceful_shutdown(&self, cluster_template: Arc<ClusterMode>) {
+        let self_clone = self.clone();
+        self_clone.request_shutdown();
+        let force_kill_time =
+            Instant::now() + Duration::from_secs_f64(cluster_template.server.shutdown_timeout);
+        while self_clone.is_alive() && force_kill_time > Instant::now() {
+            tokio::time::sleep(Duration::from_millis(100)).await;
+        }
+        if self_clone.is_alive() {
+            self_clone.force_kill();
+        }
+    }
+
+    pub(crate) fn boot_if_dead(&self, cluster_template: Arc<ClusterMode>) -> bool {
+        if !self.is_alive() {
+            if self.just_started() {
+                error!(
+                    "Worker in crash loop {:?}. Refusing to restart",
+                    self.child_pid.lock()
+                );
+                return false;
+            } else {
+                let self_clone = self.clone();
+                create_ruby_thread(move || {
+                    call_without_gvl(move || {
+                        self_clone.boot(cluster_template).ok();
+                    })
+                });
+            }
+        }
+        true
+    }
+
+    pub(crate) fn request_shutdown(&self) {
+        let child_pid = *self.child_pid.lock();
+        if let Some(pid) = child_pid {
+            if let Err(e) = kill(pid, SIGTERM) {
+                error!("Failed to send SIGTERM to process {}: {}", pid, e);
+            }
+        }
+    }
+
+    pub(crate) fn force_kill(&self) {
+        let child_pid = *self.child_pid.lock();
+        if let Some(pid) = child_pid {
+            if let Err(e) = kill(pid, SIGKILL) {
+                error!("Failed to force kill process {}: {}", pid, e);
+            }
+        }
+    }
+
+    pub(crate) fn just_started(&self) -> bool {
+        let now = Instant::now();
+        now.duration_since(self.started_at).as_millis() < 2000
+    }
+
+    pub(crate) fn is_alive(&self) -> bool {
+        let child_pid = *self.child_pid.lock();
+        if let Some(pid) = child_pid {
+            match waitpid(pid, Some(WaitPidFlag::WNOHANG)) {
+                Ok(WaitStatus::Exited(_, _)) | Ok(WaitStatus::Signaled(_, _, _)) => {
+                    return false;
+                }
+                Ok(WaitStatus::StillAlive) | Ok(_) => {}
+                Err(_) => return false,
+            }
+            match kill(pid, None) {
+                Ok(_) => true,
+                Err(errno) => !matches!(errno, Errno::ESRCH),
+            }
+        } else {
+            false
+        }
+    }
+}