itsi-scheduler 0.2.16 → 0.2.18

data/ext/itsi_server/src/server/serve_strategy/cluster_mode.rs

@@ -1,5 +1,5 @@
 use crate::ruby_types::itsi_server::itsi_server_config::ItsiServerConfig;
-use crate::server::signal::SIGNAL_HANDLER_CHANNEL;
+use crate::server::signal::{subscribe_runtime_to_signals, unsubscribe_runtime};
 use crate::server::{lifecycle_event::LifecycleEvent, process_worker::ProcessWorker};
 use itsi_error::{ItsiError, Result};
 use itsi_rb_helpers::{call_with_gvl, call_without_gvl, create_ruby_thread};
@@ -7,31 +7,32 @@ use itsi_tracing::{error, info, warn};
 use magnus::Value;
 use nix::{libc::exit, unistd::Pid};
 
+use std::sync::atomic::{AtomicBool, Ordering};
 use std::{
-    sync::{atomic::AtomicUsize, Arc},
+    sync::Arc,
     time::{Duration, Instant},
 };
 use tokio::{
     runtime::{Builder as RuntimeBuilder, Runtime},
-    sync::{broadcast, watch, Mutex},
+    sync::{watch, Mutex},
     time::{self, sleep},
 };
 use tracing::{debug, instrument};
 pub(crate) struct ClusterMode {
     pub server_config: Arc<ItsiServerConfig>,
     pub process_workers: parking_lot::Mutex<Vec<ProcessWorker>>,
-    pub lifecycle_channel: broadcast::Sender<LifecycleEvent>,
 }
 
-static WORKER_ID: AtomicUsize = AtomicUsize::new(0);
 static CHILD_SIGNAL_SENDER: parking_lot::Mutex<Option<watch::Sender<()>>> =
     parking_lot::Mutex::new(None);
 
+static RELOAD_IN_PROGRESS: AtomicBool = AtomicBool::new(false);
+
 impl ClusterMode {
     pub fn new(server_config: Arc<ItsiServerConfig>) -> Self {
         let process_workers = (0..server_config.server_params.read().workers)
-            .map(|_| ProcessWorker {
-                worker_id: WORKER_ID.fetch_add(1, std::sync::atomic::Ordering::Relaxed),
+            .map(|id| ProcessWorker {
+                worker_id: id as usize,
                 ..Default::default()
             })
             .collect();
@@ -39,7 +40,6 @@ impl ClusterMode {
         Self {
             server_config,
             process_workers: parking_lot::Mutex::new(process_workers),
-            lifecycle_channel: SIGNAL_HANDLER_CHANNEL.0.clone(),
         }
     }
 
@@ -60,6 +60,26 @@ impl ClusterMode {
         }
     }
 
+    fn next_worker_id(&self) -> usize {
+        let mut ids: Vec<usize> = self
+            .process_workers
+            .lock()
+            .iter()
+            .map(|w| w.worker_id)
+            .collect();
+        self.next_available_id_in(&mut ids)
+    }
+
+    fn next_available_id_in(&self, list: &mut [usize]) -> usize {
+        list.sort_unstable();
+        for (expected, &id) in list.iter().enumerate() {
+            if id != expected {
+                return expected;
+            }
+        }
+        list.len()
+    }
+
     #[allow(clippy::await_holding_lock)]
     pub async fn handle_lifecycle_event(
         self: Arc<Self>,
@@ -80,9 +100,11 @@ impl ClusterMode {
             LifecycleEvent::Restart => {
                 if self.server_config.check_config().await {
                     self.invoke_hook("before_restart");
+                    self.server_config.stop_watcher()?;
                     self.server_config.dup_fds()?;
                     self.shutdown().await.ok();
                     info!("Shutdown complete. Calling reload exec");
+
                     self.server_config.reload_exec()?;
                 }
                 Ok(())
@@ -91,46 +113,65 @@ impl ClusterMode {
                 if !self.server_config.check_config().await {
                     return Ok(());
                 }
+
                 let should_reexec = self.server_config.clone().reload(true)?;
+
                 if should_reexec {
+                    self.server_config.stop_watcher()?;
                     self.server_config.dup_fds()?;
                     self.shutdown().await.ok();
                     self.server_config.reload_exec()?;
                 }
-                let mut workers_to_load = self.server_config.server_params.read().workers;
-                let mut next_workers = Vec::new();
-                for worker in self.process_workers.lock().drain(..) {
-                    if workers_to_load == 0 {
-                        worker.graceful_shutdown(self.clone()).await
-                    } else {
-                        workers_to_load -= 1;
-                        worker.reboot(self.clone()).await?;
-                        next_workers.push(worker);
-                    }
+
+                if RELOAD_IN_PROGRESS
+                    .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
+                    .is_err()
+                {
+                    warn!("Reload already in progress, ignoring request");
+                    return Ok(());
                 }
-                self.process_workers.lock().extend(next_workers);
-                while workers_to_load > 0 {
-                    let mut workers = self.process_workers.lock();
+                let workers_to_load = self.server_config.server_params.read().workers;
+                let mut next_workers = Vec::new();
+                let mut old_workers = self.process_workers.lock().drain(..).collect::<Vec<_>>();
+
+                // Spawn new workers
+                for i in 0..workers_to_load {
                     let worker = ProcessWorker {
-                        worker_id: WORKER_ID.fetch_add(1, std::sync::atomic::Ordering::Relaxed),
+                        worker_id: i as usize,
                         ..Default::default()
                     };
                     let worker_clone = worker.clone();
                     let self_clone = self.clone();
-                    create_ruby_thread(move || {
-                        call_without_gvl(move || {
-                            worker_clone.boot(self_clone).ok();
-                        })
+
+                    call_with_gvl(|_| {
+                        create_ruby_thread(move || {
+                            call_without_gvl(move || match worker_clone.boot(self_clone) {
+                                Err(err) => error!("Worker boot failed {:?}", err),
+                                _ => {}
+                            })
+                        });
                     });
-                    workers.push(worker);
-                    workers_to_load -= 1
+
+                    next_workers.push(worker);
+
+                    if let Some(old) = old_workers.pop() {
+                        old.graceful_shutdown(self.clone()).await;
+                    }
+                }
+
+                for worker in old_workers {
+                    worker.graceful_shutdown(self.clone()).await;
                 }
+
+                self.process_workers.lock().extend(next_workers);
+                RELOAD_IN_PROGRESS.store(false, Ordering::SeqCst);
+
                 Ok(())
             }
             LifecycleEvent::IncreaseWorkers => {
                 let mut workers = self.process_workers.lock();
                 let worker = ProcessWorker {
-                    worker_id: WORKER_ID.fetch_add(1, std::sync::atomic::Ordering::Relaxed),
+                    worker_id: self.next_worker_id(),
                     ..Default::default()
                 };
                 let worker_clone = worker.clone();
@@ -171,6 +212,10 @@ impl ClusterMode {
                 unsafe { exit(0) };
             }
             LifecycleEvent::ChildTerminated => {
+                if RELOAD_IN_PROGRESS.load(Ordering::SeqCst) {
+                    warn!("Reload already in progress, ignoring child signal");
+                    return Ok(());
+                }
                 CHILD_SIGNAL_SENDER.lock().as_ref().inspect(|i| {
                     i.send(()).ok();
                 });
@@ -275,6 +320,7 @@ impl ClusterMode {
     pub fn run(self: Arc<Self>) -> Result<()> {
         info!("Starting in Cluster mode");
         self.invoke_hook("before_fork");
+
         self.process_workers
             .lock()
             .iter()
@@ -283,10 +329,11 @@ impl ClusterMode {
         let (sender, mut receiver) = watch::channel(());
         *CHILD_SIGNAL_SENDER.lock() = Some(sender);
 
-        let mut lifecycle_rx = self.lifecycle_channel.subscribe();
         let self_ref = self.clone();
 
         self.build_runtime().block_on(async {
+          let mut lifecycle_rx = subscribe_runtime_to_signals();
+
           let self_ref = self_ref.clone();
           let memory_check_duration = if self_ref.server_config.server_params.read().worker_memory_limit.is_some(){
             time::Duration::from_secs(15)
@@ -338,11 +385,16 @@ impl ClusterMode {
                   }
 
                 },
-                Err(e) => error!("Error receiving lifecycle_event: {:?}", e),
+                Err(e) => {
+                  debug!("Lifecycle channel closed: {:?}, exiting cluster monitor loop", e);
+                  break
+                },
               }
             }
           }
         });
+
+        unsubscribe_runtime();
         self.server_config
             .server_params
             .write()

data/ext/itsi_server/src/server/serve_strategy/single_mode.rs

@@ -4,7 +4,10 @@ use crate::{
         lifecycle_event::LifecycleEvent,
         request_job::RequestJob,
         serve_strategy::acceptor::{Acceptor, AcceptorArgs},
-        signal::{SHUTDOWN_REQUESTED, SIGNAL_HANDLER_CHANNEL},
+        signal::{
+            send_lifecycle_event, subscribe_runtime_to_signals, unsubscribe_runtime,
+            SHUTDOWN_REQUESTED,
+        },
         thread_worker::{build_thread_workers, ThreadWorker},
     },
 };
@@ -29,22 +32,20 @@ use std::{
 };
 use tokio::{
     runtime::{Builder as RuntimeBuilder, Runtime},
-    sync::{
-        broadcast,
-        watch::{self},
-    },
+    sync::watch::{self},
     task::JoinSet,
 };
 use tracing::instrument;
 
 pub struct SingleMode {
+    pub worker_id: usize,
     pub executor: Builder<TokioExecutor>,
     pub server_config: Arc<ItsiServerConfig>,
-    pub(crate) lifecycle_channel: broadcast::Sender<LifecycleEvent>,
     pub restart_requested: AtomicBool,
     pub status: RwLock<HashMap<u8, (u64, u64)>>,
 }
 
+#[derive(PartialEq, Debug)]
 pub enum RunningPhase {
     Running,
     ShutdownPending,
@@ -53,31 +54,45 @@ pub enum RunningPhase {
 
 impl SingleMode {
     #[instrument(parent=None, skip_all)]
-    pub fn new(server_config: Arc<ItsiServerConfig>) -> Result<Self> {
+    pub fn new(server_config: Arc<ItsiServerConfig>, worker_id: usize) -> Result<Self> {
         server_config.server_params.read().preload_ruby()?;
-        let mut executor = Builder::new(TokioExecutor::new());
-        executor
-            .http1()
-            .header_read_timeout(server_config.server_params.read().header_read_timeout)
-            .writev(true)
-            .timer(TokioTimer::new());
-        executor
-            .http2()
-            .max_concurrent_streams(100)
-            .max_local_error_reset_streams(100)
-            .enable_connect_protocol()
-            .max_header_list_size(10 * 1024 * 1024)
-            .max_send_buf_size(16 * 1024 * 1024);
+        let executor = {
+            let mut executor = Builder::new(TokioExecutor::new());
+            let server_params = server_config.server_params.read();
+            let mut http1_executor = executor.http1();
+
+            http1_executor
+                .header_read_timeout(server_params.header_read_timeout)
+                .pipeline_flush(server_params.pipeline_flush)
+                .timer(TokioTimer::new());
+
+            if let Some(writev) = server_params.writev {
+                http1_executor.writev(writev);
+            }
+
+            executor
+                .http2()
+                .max_concurrent_streams(server_params.max_concurrent_streams)
+                .max_local_error_reset_streams(server_params.max_local_error_reset_streams)
+                .max_header_list_size(server_params.max_header_list_size)
+                .max_send_buf_size(server_params.max_send_buf_size)
+                .enable_connect_protocol();
+            executor
+        };
 
         Ok(Self {
+            worker_id,
             executor,
             server_config,
-            lifecycle_channel: SIGNAL_HANDLER_CHANNEL.0.clone(),
             restart_requested: AtomicBool::new(false),
             status: RwLock::new(HashMap::new()),
         })
     }
 
+    pub fn is_zero_worker(&self) -> bool {
+        self.worker_id == 0
+    }
+
     pub fn build_runtime(&self) -> Runtime {
         let mut builder: RuntimeBuilder = if self
             .server_config
@@ -103,7 +118,7 @@ impl SingleMode {
 
     pub fn stop(&self) -> Result<()> {
         SHUTDOWN_REQUESTED.store(true, std::sync::atomic::Ordering::SeqCst);
-        self.lifecycle_channel.send(LifecycleEvent::Shutdown).ok();
+        send_lifecycle_event(LifecycleEvent::Shutdown);
         Ok(())
     }
 
@@ -182,7 +197,7 @@ impl SingleMode {
                         .unwrap();
                     let receiver = self.clone();
                     monitor_runtime.block_on({
-                        let mut lifecycle_rx = receiver.lifecycle_channel.subscribe();
+                        let mut lifecycle_rx = subscribe_runtime_to_signals();
                         let receiver = receiver.clone();
                         let thread_workers = thread_workers.clone();
                         async move {
@@ -201,18 +216,19 @@ impl SingleMode {
                                   }
                                   lifecycle_event = lifecycle_rx.recv() => {
                                       match lifecycle_event {
-                                          Ok(LifecycleEvent::Restart) => {
+                                          Ok(LifecycleEvent::Restart) | Ok(LifecycleEvent::Reload) => {
                                               receiver.restart().await.ok();
                                           }
-                                          Ok(LifecycleEvent::Reload) => {
-                                              receiver.reload().await.ok();
-                                          }
                                           Ok(LifecycleEvent::Shutdown) => {
                                             break;
                                           }
                                           Ok(LifecycleEvent::PrintInfo) => {
                                             receiver.print_info(thread_workers.clone()).await.ok();
                                           }
+                                          Err(e) => {
+                                            debug!("Lifecycle channel closed: {:?}, exiting single mode monitor loop", e);
+                                            break;
+                                          }
                                           _ => {}
                                       }
                                   }
@@ -227,13 +243,15 @@ impl SingleMode {
 
     #[instrument(name="worker", parent=None, skip(self), fields(pid=format!("{}", Pid::this())))]
     pub fn run(self: Arc<Self>) -> Result<()> {
-        let (thread_workers, job_sender, nonblocking_sender) =
-            build_thread_workers(self.server_config.server_params.read().clone(), Pid::this())
-                .inspect_err(|e| {
-                    if let Some(err_val) = e.value() {
-                        print_rb_backtrace(err_val);
-                    }
-                })?;
+        let (thread_workers, job_sender, nonblocking_sender) = build_thread_workers(
+            self.server_config.server_params.read().clone(),
+            self.worker_id,
+        )
+        .inspect_err(|e| {
+            if let Some(err_val) = e.value() {
+                print_rb_backtrace(err_val);
+            }
+        })?;
 
         let worker_count = thread_workers.len();
         info!(
@@ -244,6 +262,14 @@ impl SingleMode {
         let shutdown_timeout = self.server_config.server_params.read().shutdown_timeout;
         let (shutdown_sender, _) = watch::channel(RunningPhase::Running);
         let monitor_thread = self.clone().start_monitors(thread_workers.clone());
+
+        // If we're on Linux with reuse_port enabled, we can use
+        // kernel level load balancing across processes sharing a port.
+        // To take advantage of this, these forks will rebind to the same port upon boot.
+        // Worker 0 is special (this one just inherits the bind from the master process).
+        let is_zero_worker = self.is_zero_worker();
+        let should_rebind = !is_zero_worker && self.server_config.use_reuse_port_load_balancing();
+
         if monitor_thread.is_none() {
             error!("Failed to start monitor thread");
             return Err(ItsiError::new("Failed to start monitor thread"));
@@ -253,106 +279,123 @@ impl SingleMode {
             return Ok(());
         }
         let runtime = self.build_runtime();
-        let result = runtime.block_on(
-          async  {
-              let mut listener_task_set = JoinSet::new();
-              let server_params = self.server_config.server_params.read().clone();
-              if let Err(err) = server_params.initialize_middleware().await {
-                  error!("Failed to initialize middleware: {}", err);
-                  return Err(ItsiError::new("Failed to initialize middleware"))
-              }
-              let tokio_listeners = server_params.listeners.lock()
-                  .drain(..)
-                  .map(|list| {
-                    Arc::new(list.into_tokio_listener())
-                  })
-                  .collect::<Vec<_>>();
-
-              tokio_listeners.iter().cloned().for_each(|listener| {
-                  let shutdown_sender = shutdown_sender.clone();
-                  let job_sender = job_sender.clone();
-                  let nonblocking_sender = nonblocking_sender.clone();
-
-                  let mut lifecycle_rx = self.lifecycle_channel.subscribe();
-                  let mut shutdown_receiver = shutdown_sender.subscribe();
-                  let mut acceptor = Acceptor{
-                      acceptor_args: Arc::new(
-                        AcceptorArgs{
-                          strategy: self.clone(),
-                          listener_info: listener.listener_info(),
-                          shutdown_receiver: shutdown_sender.subscribe(),
-                          job_sender: job_sender.clone(),
-                          nonblocking_sender: nonblocking_sender.clone(),
-                          server_params: server_params.clone()
-                        }
-                      ),
-                      join_set: JoinSet::new()
-                  };
-
-                  let shutdown_rx_for_acme_task = shutdown_receiver.clone();
-                  let acme_task_listener_clone = listener.clone();
-                  listener_task_set.spawn(async move {
-                      acme_task_listener_clone.spawn_acme_event_task(shutdown_rx_for_acme_task).await;
-                  });
-
-                  listener_task_set.spawn(async move {
-                      loop {
-                          tokio::select! {
-                              accept_result = listener.accept() => {
-                                  match accept_result {
-                                      Ok(accepted) => acceptor.serve_connection(accepted).await,
-                                      Err(e) => debug!("Listener.accept failed: {:?}", e)
-                                  }
-                              },
-                              _ = shutdown_receiver.changed() => {
-                                  debug!("Shutdown requested via receiver");
-                                  break;
-                              },
-                              lifecycle_event = lifecycle_rx.recv() => {
-                                  match lifecycle_event {
-                                      Ok(LifecycleEvent::Shutdown) => {
-                                          debug!("Received LifecycleEvent::Shutdown");
-                                          let _ = shutdown_sender.send(RunningPhase::ShutdownPending);
-                                          for _ in 0..worker_count {
-                                              let _ = job_sender.send(RequestJob::Shutdown).await;
-                                              let _ = nonblocking_sender.send(RequestJob::Shutdown).await;
-                                          }
-                                          break;
-                                      },
-                                      Err(e) =>  error!("Error receiving lifecycle event: {:?}", e),
-                                      _ => ()
+        let result = runtime.block_on(async {
+            let mut listener_task_set = JoinSet::new();
+            let server_params = self.server_config.server_params.read().clone();
+            if let Err(err) = server_params.initialize_middleware().await {
+                error!("Failed to initialize middleware: {}", err);
+                return Err(ItsiError::new("Failed to initialize middleware"));
+            }
+            let tokio_listeners = server_params
+                .listeners
+                .lock()
+                .drain(..)
+                .map(|list| Arc::new(list.into_tokio_listener(should_rebind)))
+                .collect::<Vec<_>>();
+
+            tokio_listeners.iter().cloned().for_each(|listener| {
+                let shutdown_sender = shutdown_sender.clone();
+                let job_sender = job_sender.clone();
+                let nonblocking_sender = nonblocking_sender.clone();
+
+                let mut lifecycle_rx = subscribe_runtime_to_signals();
+                let mut shutdown_receiver = shutdown_sender.subscribe();
+                let mut acceptor = Acceptor {
+                    acceptor_args: Arc::new(AcceptorArgs {
+                        strategy: self.clone(),
+                        listener_info: listener.listener_info(),
+                        shutdown_receiver: shutdown_sender.subscribe(),
+                        job_sender: job_sender.clone(),
+                        nonblocking_sender: nonblocking_sender.clone(),
+                        server_params: server_params.clone(),
+                    }),
+                    join_set: JoinSet::new(),
+                };
+
+                let shutdown_rx_for_acme_task = shutdown_receiver.clone();
+                let acme_task_listener_clone = listener.clone();
+
+                let mut after_accept_wait: Option<Duration> = None::<Duration>;
+
+                if cfg!(target_os = "macos") {
+                    after_accept_wait = if server_params.workers > 1 && !(server_params.socket_opts.reuse_port && server_params.socket_opts.reuse_address) {
+                        Some(Duration::from_nanos(10 * server_params.workers as u64))
+                    } else {
+                        None
+                    };
+                };
+
+                listener_task_set.spawn(async move {
+                    acme_task_listener_clone
+                        .spawn_acme_event_task(shutdown_rx_for_acme_task)
+                        .await;
+                });
+
+                listener_task_set.spawn(async move {
+                    loop {
+                        // Process any pending signals before select
+                        tokio::select! {
+                            accept_result = listener.accept() => {
+                                match accept_result {
+                                    Ok(accepted) => acceptor.serve_connection(accepted).await,
+                                    Err(e) => debug!("Listener.accept failed: {:?}", e)
+                                }
+                                if cfg!(target_os = "macos") {
+                                  if let Some(after_accept_wait) = after_accept_wait{
+                                    tokio::time::sleep(after_accept_wait).await;
                                   }
-                              }
-                          }
-                      }
-                      acceptor.join().await;
-                  });
-              });
-
-              if self.is_single_mode() {
+                                }
+                            },
+                            _ = shutdown_receiver.changed() => {
+                                debug!("Shutdown requested via receiver");
+                                break;
+                            },
+                            lifecycle_event = lifecycle_rx.recv() => {
+                                match lifecycle_event {
+                                    Ok(LifecycleEvent::Shutdown) => {
+                                        debug!("Received LifecycleEvent::Shutdown");
+                                        let _ = shutdown_sender.send(RunningPhase::ShutdownPending);
+                                        break;
+                                    },
+                                    Err(e) => {
+                                      debug!("Lifecycle channel closed: {:?}, exiting accept loop", e);
+                                      break
+                                    },
+                                    _ => ()
+                                }
+                            }
+                        }
+                    }
+                    acceptor.join().await;
+                });
+            });
+
+            if self.is_single_mode() {
                 self.invoke_hook("after_start");
-              }
+            }
 
-              while let Some(_res) = listener_task_set.join_next().await {}
-              drop(tokio_listeners);
+            while let Some(_res) = listener_task_set.join_next().await {}
+            drop(tokio_listeners);
 
-              Ok::<(), ItsiError>(())
-          });
+            Ok::<(), ItsiError>(())
+        });
 
         debug!("Single mode runtime exited.");
 
+        for _i in 0..thread_workers.len() {
+            job_sender.send_blocking(RequestJob::Shutdown).unwrap();
+            nonblocking_sender
+                .send_blocking(RequestJob::Shutdown)
+                .unwrap();
+        }
         if result.is_err() {
-            for _i in 0..thread_workers.len() {
-                job_sender.send_blocking(RequestJob::Shutdown).unwrap();
-                nonblocking_sender
-                    .send_blocking(RequestJob::Shutdown)
-                    .unwrap();
-            }
-            self.lifecycle_channel.send(LifecycleEvent::Shutdown).ok();
+            send_lifecycle_event(LifecycleEvent::Shutdown);
         }
 
         shutdown_sender.send(RunningPhase::Shutdown).ok();
         runtime.shutdown_timeout(Duration::from_millis(100));
+        unsubscribe_runtime();
+
         debug!("Shutdown timeout finished.");
 
         let deadline = Instant::now() + Duration::from_secs_f64(shutdown_timeout);
@@ -384,26 +427,6 @@ impl SingleMode {
     pub fn is_single_mode(&self) -> bool {
         self.server_config.server_params.read().workers == 1
     }
-    /// Attempts to reload the config "live"
-    /// Not that when running in single mode this will not unload
-    /// old code. If you need a clean restart, use the `restart` (SIGHUP) method instead
-    pub async fn reload(&self) -> Result<()> {
-        if !self.server_config.check_config().await {
-            return Ok(());
-        }
-        let should_reexec = self.server_config.clone().reload(false)?;
-        if should_reexec {
-            if self.is_single_mode() {
-                self.invoke_hook("before_restart");
-            }
-            self.server_config.dup_fds()?;
-            self.server_config.reload_exec()?;
-        }
-        self.restart_requested.store(true, Ordering::SeqCst);
-        self.stop()?;
-        self.server_config.server_params.read().preload_ruby()?;
-        Ok(())
-    }
 
     pub fn invoke_hook(&self, hook_name: &str) {
         if let Some(hook) = self.server_config.server_params.read().hooks.get(hook_name) {
@@ -418,6 +441,7 @@ impl SingleMode {
         if self.is_single_mode() {
             self.invoke_hook("before_restart");
         }
+        self.server_config.stop_watcher()?;
         self.server_config.dup_fds()?;
         self.server_config.reload_exec()?;
         Ok(())