Allow setting the number of runtime workers to be used. (#258)

This change adds a new configuration parameter called `worker_threads` that
allows setting the number of workers the Tokio Runtime will use. It defaults to
4 to maintain backward compatibility.

Given that the config file parse is done asynchronously, first, a transient runtime
is created for reading config, and once it has been parsed, the actual runtime
that will be used for PgCat execution is created.

pgcat.toml

@@ -44,6 +44,9 @@ log_client_disconnections = false
 # Reload config automatically if it changes.
 autoreload = false
 
+# Number of worker threads the Runtime will use (4 by default).
+worker_threads = 5
+
 # TLS
 # tls_certificate = "server.cert"
 # tls_private_key = "server.key"

src/config.rs

@@ -178,6 +178,9 @@ pub struct General {
     #[serde(default = "General::default_ban_time")]
     pub ban_time: i64,
 
+    #[serde(default = "General::default_worker_threads")]
+    pub worker_threads: usize,
+
     #[serde(default)] // False
     pub autoreload: bool,
 
@@ -219,6 +222,10 @@ impl General {
     pub fn default_ban_time() -> i64 {
         60
     }
+
+    pub fn default_worker_threads() -> usize {
+        4
+    }
 }
 
 impl Default for General {
@@ -234,6 +241,7 @@ impl Default for General {
             healthcheck_timeout: Self::default_healthcheck_timeout(),
             healthcheck_delay: Self::default_healthcheck_delay(),
             ban_time: Self::default_ban_time(),
+            worker_threads: Self::default_worker_threads(),
             log_client_connections: false,
             log_client_disconnections: false,
             autoreload: false,

src/main.rs

@@ -49,6 +49,7 @@ use parking_lot::Mutex;
 use pgcat::format_duration;
 use tokio::net::TcpListener;
 use tokio::{
+    runtime::Builder,
     signal::unix::{signal as unix_signal, SignalKind},
     sync::mpsc,
 };
@@ -79,8 +80,7 @@ use crate::pool::{ClientServerMap, ConnectionPool};
 use crate::prometheus::start_metric_server;
 use crate::stats::{Collector, Reporter, REPORTER};
 
-#[tokio::main(worker_threads = 4)]
-async fn main() {
+fn main() -> Result<(), Box<dyn std::error::Error>> {
     env_logger::builder().format_timestamp_micros().init();
 
     info!("Welcome to PgCat! Meow. (Version {})", VERSION);
@@ -98,6 +98,11 @@ async fn main() {
         String::from("pgcat.toml")
     };
 
+    // Create a transient runtime for loading the config for the first time.
+    {
+        let runtime = Builder::new_multi_thread().worker_threads(1).build()?;
+
+        runtime.block_on(async {
             match config::parse(&config_file).await {
                 Ok(_) => (),
                 Err(err) => {
@@ -105,9 +110,19 @@ async fn main() {
                     std::process::exit(exitcode::CONFIG);
                 }
             };
+        });
+    }
 
     let config = get_config();
 
+    // Create the runtime now we know required worker_threads.
+    let runtime = Builder::new_multi_thread()
+        .worker_threads(config.general.worker_threads)
+        .enable_all()
+        .build()?;
+
+    runtime.block_on(async move {
+
 	if let Some(true) = config.general.enable_prometheus_exporter {
             let http_addr_str = format!(
 		"{}:{}",
@@ -309,4 +324,6 @@ async fn main() {
 	}
 
 	info!("Shutting down...");
+    });
+    Ok(())
 }