Allow manually specifying the number of threads

Cargo.toml

@@ -4,7 +4,7 @@ description = "A very simple secrets vault."
 authors = ["Dennis Schubert <[email protected]>"]
 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",

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.

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://[email protected]/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)]

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<PgPool, sqlx::Error> {
     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::<SocketAddr>(),

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://[email protected]/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<usize>,
+}