Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 1 commit into from
Dec 26, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions bin/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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 }
Expand Down
20 changes: 11 additions & 9 deletions bin/src/ctl/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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> {
FlorentinDUBOIS marked this conversation as resolved.
Show resolved Hide resolved
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 } => {
Expand Down
22 changes: 22 additions & 0 deletions bin/src/lib.rs
Original file line number Diff line number Diff line change
@@ -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;
60 changes: 1 addition & 59 deletions bin/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -134,61 +131,6 @@ fn main(args: Args) -> Result<(), MainError> {
result
}

/// Set workers process affinity, see man sched_setaffinity
FlorentinDUBOIS marked this conversation as resolved.
Show resolved Hide resolved
/// 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
Expand Down
60 changes: 59 additions & 1 deletion bin/src/util.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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);
};
}
Loading