Merge pull request #1165 from anonkey/allow-using-sozu-as-lib

feat: allow sozu to be used as lib by another package

bin/Cargo.toml

@@ -18,6 +18,14 @@ edition = "2021"
 rust-version = "1.80.0"
 include = ["README.md", "Cargo.toml", "src/**/*"]
 
+[lib]
+name = "sozu"
+path = "src/lib.rs"
+
+[[bin]]
+name = "sozu"
+path = "src/main.rs"
+
 [dependencies]
 clap = { version = "^4.5.20", features = ["derive"] }
 jemallocator = { version = "^0.5.4", optional = true }

bin/src/ctl/mod.rs

@@ -84,25 +84,27 @@ pub fn ctl(args: cli::Args) -> Result<(), CtlError> {
         std::process::exit(0);
     }
 
-    let channel = create_channel(&config)?;
-
     let timeout = Duration::from_millis(args.timeout.unwrap_or(config.ctl_command_timeout));
     if !args.json {
         debug!("applying timeout {:?}", timeout);
     }
 
-    let mut command_manager = CommandManager {
-        channel,
-        timeout,
-        config,
-        json: args.json,
-    };
+    let mut command_manager = CommandManager::new(config, timeout, args.json)?;
 
     command_manager.handle_command(args.cmd)
 }
 
 impl CommandManager {
-    fn handle_command(&mut self, command: SubCmd) -> Result<(), CtlError> {
+    pub fn new(config: Config, timeout: Duration, json: bool) -> Result<Self, CtlError> {
+        Ok(Self {
+            channel: create_channel(&config)?,
+            timeout,
+            config,
+            json,
+        })
+    }
+
+    pub fn handle_command(&mut self, command: SubCmd) -> Result<(), CtlError> {
         debug!("Executing command {:?}", command);
         match command {
             SubCmd::Shutdown { hard } => {

bin/src/lib.rs

@@ -0,0 +1,22 @@
+#[macro_use]
+extern crate sozu_lib as sozu;
+#[macro_use]
+extern crate sozu_command_lib;
+
+#[cfg(feature = "jemallocator")]
+#[global_allocator]
+static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc;
+
+/// the arguments to the sozu command line
+pub mod cli;
+/// Receives orders from the CLI, transmits to workers
+// mod command;
+pub mod command;
+/// The command line logic
+pub mod ctl;
+/// Forking & restarting the main process using a more recent executable of Sōzu
+mod upgrade;
+/// Some unix helper functions
+pub mod util;
+/// Start and restart the worker UNIX processes
+mod worker;

bin/src/main.rs

@@ -50,13 +50,10 @@ mod worker;
 
 use std::panic;
 
-#[cfg(target_os = "linux")]
-use libc::{cpu_set_t, pid_t};
-
 use sozu::metrics::METRICS;
 
 use cli::Args;
-use command::{begin_main_process, sessions::WorkerSession, StartError};
+use command::{begin_main_process, StartError};
 use ctl::CtlError;
 use upgrade::UpgradeError;
 use worker::WorkerError;
@@ -134,61 +131,6 @@ fn main(args: Args) -> Result<(), MainError> {
     result
 }
 
-/// Set workers process affinity, see man sched_setaffinity
-/// Bind each worker (including the main) process to a CPU core.
-/// Can bind multiple processes to a CPU core if there are more processes
-/// than CPU cores. Only works on Linux.
-#[cfg(target_os = "linux")]
-fn set_workers_affinity(workers: &Vec<WorkerSession>) {
-    let mut cpu_count = 0;
-    let max_cpu = num_cpus::get();
-
-    // +1 for the main process that will also be bound to its CPU core
-    if (workers.len() + 1) > max_cpu {
-        warn!(
-            "There are more workers than available CPU cores, \
-          multiple workers will be bound to the same CPU core. \
-          This may impact performances"
-        );
-    }
-
-    let main_pid = unsafe { libc::getpid() };
-    set_process_affinity(main_pid, cpu_count);
-    cpu_count += 1;
-
-    for worker in workers {
-        if cpu_count >= max_cpu {
-            cpu_count = 0;
-        }
-
-        set_process_affinity(worker.pid, cpu_count);
-
-        cpu_count += 1;
-    }
-}
-
-/// Set workers process affinity, see man sched_setaffinity
-/// Bind each worker (including the main) process to a CPU core.
-/// Can bind multiple processes to a CPU core if there are more processes
-/// than CPU cores. Only works on Linux.
-#[cfg(not(target_os = "linux"))]
-fn set_workers_affinity(_: &Vec<Worker>) {}
-
-/// Set a specific process to run onto a specific CPU core
-#[cfg(target_os = "linux")]
-use std::mem;
-#[cfg(target_os = "linux")]
-fn set_process_affinity(pid: pid_t, cpu: usize) {
-    unsafe {
-        let mut cpu_set: cpu_set_t = mem::zeroed();
-        let size_cpu_set = mem::size_of::<cpu_set_t>();
-        libc::CPU_SET(cpu, &mut cpu_set);
-        libc::sched_setaffinity(pid, size_cpu_set, &cpu_set);
-
-        debug!("Worker {} bound to CPU core {}", pid, cpu);
-    };
-}
-
 fn register_panic_hook() {
     // We save the original panic hook so we can call it later
     // to have the original behavior

bin/src/util.rs

@@ -14,7 +14,10 @@ use nix::{
 use sozu_command_lib::config::Config;
 use sozu_lib::metrics::{self, MetricError};
 
-use crate::cli;
+use crate::{cli, command};
+
+#[cfg(target_os = "linux")]
+use libc::{cpu_set_t, pid_t};
 
 #[derive(thiserror::Error, Debug)]
 pub enum UtilError {
@@ -173,3 +176,58 @@ pub unsafe fn get_executable_path() -> Result<String, UtilError> {
 
     Ok(path.to_string_lossy().to_string())
 }
+
+/// Set workers process affinity, see man sched_setaffinity
+/// Bind each worker (including the main) process to a CPU core.
+/// Can bind multiple processes to a CPU core if there are more processes
+/// than CPU cores. Only works on Linux.
+#[cfg(target_os = "linux")]
+pub fn set_workers_affinity(workers: &Vec<command::sessions::WorkerSession>) {
+    let mut cpu_count = 0;
+    let max_cpu = num_cpus::get();
+
+    // +1 for the main process that will also be bound to its CPU core
+    if (workers.len() + 1) > max_cpu {
+        warn!(
+            "There are more workers than available CPU cores, \
+          multiple workers will be bound to the same CPU core. \
+          This may impact performances"
+        );
+    }
+
+    let main_pid = unsafe { libc::getpid() };
+    set_process_affinity(main_pid, cpu_count);
+    cpu_count += 1;
+
+    for worker in workers {
+        if cpu_count >= max_cpu {
+            cpu_count = 0;
+        }
+
+        set_process_affinity(worker.pid, cpu_count);
+
+        cpu_count += 1;
+    }
+}
+
+/// Set workers process affinity, see man sched_setaffinity
+/// Bind each worker (including the main) process to a CPU core.
+/// Can bind multiple processes to a CPU core if there are more processes
+/// than CPU cores. Only works on Linux.
+#[cfg(not(target_os = "linux"))]
+pub fn set_workers_affinity(_: &Vec<command::sessions::WorkerSession>) {}
+
+/// Set a specific process to run onto a specific CPU core
+#[cfg(target_os = "linux")]
+use std::mem;
+#[cfg(target_os = "linux")]
+pub fn set_process_affinity(pid: pid_t, cpu: usize) {
+    unsafe {
+        let mut cpu_set: cpu_set_t = mem::zeroed();
+        let size_cpu_set = mem::size_of::<cpu_set_t>();
+        libc::CPU_SET(cpu, &mut cpu_set);
+        libc::sched_setaffinity(pid, size_cpu_set, &cpu_set);
+
+        debug!("Worker {} bound to CPU core {}", pid, cpu);
+    };
+}