From f6e4056183ddfb3537a17459023a7b6bf0e38089 Mon Sep 17 00:00:00 2001 From: Dennis Schubert Date: Sun, 5 Jan 2025 09:10:26 +0100 Subject: [PATCH] Allow manually specifying the number of threads --- Cargo.lock | 19 +------------------ Cargo.toml | 3 +-- Changelog.md | 4 ++++ src/lib.rs | 18 ++---------------- src/main.rs | 33 +++++++++++++++++++++++---------- src/settings.rs | 20 ++++++++++++++++++++ 6 files changed, 51 insertions(+), 46 deletions(-) create mode 100644 src/settings.rs diff --git a/Cargo.lock b/Cargo.lock index f005fae..89479e4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -684,12 +684,6 @@ version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" -[[package]] -name = "hermit-abi" -version = "0.3.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d231dfb89cfffdbc30e7fc41579ed6066ad03abda9e567ccafae602b97ec5024" - [[package]] name = "hex" version = "0.4.3" @@ -1171,16 +1165,6 @@ dependencies = [ "libm", ] -[[package]] -name = "num_cpus" -version = "1.16.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4161fcb6d602d4d2081af7c3a45852d875a03dd337a6bfdd6e06407b61342a43" -dependencies = [ - "hermit-abi", - "libc", -] - [[package]] name = "object" version = "0.36.7" @@ -2106,14 +2090,13 @@ checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" [[package]] name = "vssv" -version = "1.1.3" +version = "1.2.0" dependencies = [ "anyhow", "axum", "axum-extra", "chrono", "clap", - "num_cpus", "serde", "sqlx", "thiserror", diff --git a/Cargo.toml b/Cargo.toml index dbb5deb..b69a810 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -4,7 +4,7 @@ description = "A very simple secrets vault." authors = ["Dennis Schubert "] repository = "https://github.com/denschub/vssv" license = "MIT" -version = "1.1.3" +version = "1.2.0" edition = "2021" [profile.release] @@ -16,7 +16,6 @@ axum = { version = "0.8", features = ["macros"] } axum-extra = { version = "0.10", features = ["typed-header"] } chrono = "0.4" clap = { version = "4", features = ["derive", "env"] } -num_cpus = "1" serde = { version = "1", features = ["derive"] } sqlx = { version = "0.8", features = [ "chrono", diff --git a/Changelog.md b/Changelog.md index ec3ab87..e77eab9 100644 --- a/Changelog.md +++ b/Changelog.md @@ -1,3 +1,7 @@ +# 1.2.0 + +This version introduces a new setting, `--threads`/`THREADS` that allows limiting the number of worker threads and the size of the database connection pool. If this flag is not set, the number of available CPU cores will be used, which matches the current behavior. + # 1.1.3 This version does not contain any functional changes. It only updates third-party dependencies. diff --git a/src/lib.rs b/src/lib.rs index 09a958d..ff99150 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,23 +1,9 @@ -use std::net::SocketAddr; - -use sqlx::{postgres::PgConnectOptions, PgPool}; +use sqlx::PgPool; pub mod errors; pub mod models; pub mod routes; - -#[derive(Debug, clap::Parser)] -#[clap(about, version, propagate_version = true)] -pub struct Cli { - /// The database URL to connect to. Needs to be a valid PostgreSQL - /// connection URL, like `postgres://postgres@127.0.0.1/vssv` - #[clap(long, short, env = "DATABASE_URL")] - pub database_url: PgConnectOptions, - - /// The Socket Address the server should listen on - #[clap(long, short, env = "LISTEN_ADDR", default_value = "[::1]:3000")] - pub listen_addr: SocketAddr, -} +pub mod settings; /// Holds the web server's state. #[derive(Clone)] diff --git a/src/main.rs b/src/main.rs index c468775..504e08b 100644 --- a/src/main.rs +++ b/src/main.rs @@ -8,7 +8,7 @@ use sqlx::{ use tokio::{net::TcpListener, signal}; use tracing::info; -use vssv::{routes::build_router, Cli, ServerState}; +use vssv::{routes::build_router, settings::Settings, ServerState}; /// Listens to SIGINT (aka ctrl-c) and SIGTERM and completes whenever one of /// those signals happen. @@ -34,30 +34,43 @@ async fn shutdown_signal() { } /// Creates a [PgPool] if possible. The pool has its max_connections value set -/// to the number of CPUs available. +/// to mirror the tokio worker thread count. pub async fn get_db_pool(connect_options: PgConnectOptions) -> Result { PgPoolOptions::new() .max_connections( - num_cpus::get() + tokio::runtime::Handle::current() + .metrics() + .num_workers() .try_into() - .expect("number of CPU cores should fit into an u32"), + .expect("num_workers to be less than 2^32"), ) .connect_with(connect_options) .await } -#[tokio::main] -async fn main() -> anyhow::Result<()> { - let cli = Cli::parse(); +fn main() -> anyhow::Result<()> { + let settings = Settings::parse(); + + let mut rt = tokio::runtime::Builder::new_multi_thread(); + if let Some(threads) = settings.threads { + rt.worker_threads(threads); + } + + rt.enable_all() + .build()? + .block_on(async { run(settings).await }) +} + +async fn run(settings: Settings) -> anyhow::Result<()> { tracing_subscriber::fmt::init(); - let db_pool = get_db_pool(cli.database_url).await?; + let db_pool = get_db_pool(settings.database_url).await?; sqlx::migrate!().run(&db_pool).await?; let router = build_router(ServerState { database: db_pool }); - let listener = TcpListener::bind(cli.listen_addr).await?; + let listener = TcpListener::bind(settings.listen_addr).await?; - info!("Will start to listen on `{}`...", cli.listen_addr); + info!("Will start to listen on `{}`...", settings.listen_addr); axum::serve( listener, router.into_make_service_with_connect_info::(), diff --git a/src/settings.rs b/src/settings.rs new file mode 100644 index 0000000..856405a --- /dev/null +++ b/src/settings.rs @@ -0,0 +1,20 @@ +use std::net::SocketAddr; + +use sqlx::postgres::PgConnectOptions; + +#[derive(Debug, clap::Parser)] +#[clap(about, version, propagate_version = true)] +pub struct Settings { + /// The database URL to connect to. Needs to be a valid PostgreSQL + /// connection URL, like `postgres://postgres@127.0.0.1/vssv` + #[clap(long, short, env = "DATABASE_URL")] + pub database_url: PgConnectOptions, + + /// The Socket Address the server should listen on + #[clap(long, short, env = "LISTEN_ADDR", default_value = "[::1]:3000")] + pub listen_addr: SocketAddr, + + /// Limits the number of threads used - defaults to the number of CPU cores + #[clap(long, env = "THREADS")] + pub threads: Option, +}