diff --git a/pgdog/src/backend/pool/mod.rs b/pgdog/src/backend/pool/mod.rs index 2da7c94..a230556 100644 --- a/pgdog/src/backend/pool/mod.rs +++ b/pgdog/src/backend/pool/mod.rs @@ -22,7 +22,7 @@ pub use connection::Connection; pub use error::Error; pub use guard::Guard; pub use healthcheck::Healtcheck; -pub use monitor::Monitor; +use monitor::Monitor; pub use pool::Pool; pub use replicas::Replicas; pub use shard::Shard; diff --git a/pgdog/src/backend/pool/monitor.rs b/pgdog/src/backend/pool/monitor.rs index 153c065..b6e33c1 100644 --- a/pgdog/src/backend/pool/monitor.rs +++ b/pgdog/src/backend/pool/monitor.rs @@ -44,13 +44,18 @@ use tracing::info; use tracing::{debug, error}; /// Pool maintenance. -pub struct Monitor { +/// +/// See [`crate::backend::pool::monitor`] module documentation +/// for more details. +pub(super) struct Monitor { pool: Pool, } impl Monitor { /// Launch the pool maintenance loop. - pub fn new(pool: &Pool) { + /// + /// This is done automatically when the pool is created. + pub(super) fn new(pool: &Pool) { let monitor = Self { pool: pool.clone() }; spawn(async move { @@ -140,6 +145,9 @@ impl Monitor { debug!("maintenance loop is shut down [{}]", self.pool.addr()); } + /// The healthcheck loop. + /// + /// Runs regularly and ensures the pool triggers healthchecks on idle connections. async fn healthchecks(pool: Pool) { let mut tick = interval(pool.lock().config().idle_healthcheck_interval()); let comms = pool.comms(); diff --git a/pgdog/src/backend/pool/replicas.rs b/pgdog/src/backend/pool/replicas.rs index a94fdeb..d7d8188 100644 --- a/pgdog/src/backend/pool/replicas.rs +++ b/pgdog/src/backend/pool/replicas.rs @@ -79,17 +79,17 @@ impl Replicas { let mut candidates = self .pools .iter() - .map(|pool| (pool.state(), pool)) + .map(|pool| (pool.banned(), pool)) .collect::>(); if let Some(primary) = primary { - candidates.push((primary.state(), primary)); + candidates.push((primary.banned(), primary)); } candidates.shuffle(&mut rand::thread_rng()); // All replicas are banned, unban everyone. - let banned = candidates.iter().all(|(state, _)| state.banned); + let banned = candidates.iter().all(|(banned, _)| *banned); if banned { candidates .iter() diff --git a/pgdog/src/cli.rs b/pgdog/src/cli.rs index 4e6b876..c4ad2e6 100644 --- a/pgdog/src/cli.rs +++ b/pgdog/src/cli.rs @@ -7,9 +7,9 @@ use clap::Parser; #[derive(Parser, Debug)] pub struct Cli { /// Path to the configuration file. Default: "pgdog.toml" - #[clap(default_value = "pgdog.toml")] - config: PathBuf, + #[arg(default_value = "pgdog.toml")] + pub config: PathBuf, /// Path to the users.toml file. Default: "users.toml" - #[clap(default_value = "users.toml")] - users: PathBuf, + #[arg(default_value = "users.toml")] + pub users: PathBuf, } diff --git a/pgdog/src/config/mod.rs b/pgdog/src/config/mod.rs index 564c20f..9aad7e0 100644 --- a/pgdog/src/config/mod.rs +++ b/pgdog/src/config/mod.rs @@ -4,9 +4,9 @@ pub mod error; use error::Error; -use std::collections::HashMap; use std::fs::read_to_string; use std::sync::Arc; +use std::{collections::HashMap, path::PathBuf}; use arc_swap::ArcSwap; use once_cell::sync::Lazy; @@ -22,8 +22,8 @@ pub fn config() -> Arc { } /// Load the configuration file from disk. -pub fn load() -> Result { - let config = ConfigAndUsers::load()?; +pub fn load(config: &PathBuf, users: &PathBuf) -> Result { + let config = ConfigAndUsers::load(config, users)?; CONFIG.store(Arc::new(config.clone())); Ok(config) } @@ -38,8 +38,8 @@ pub struct ConfigAndUsers { impl ConfigAndUsers { /// Load configuration from disk or use defaults. - pub fn load() -> Result { - let config: Config = if let Ok(config) = read_to_string("pgdog.toml") { + pub fn load(config: &PathBuf, users: &PathBuf) -> Result { + let config: Config = if let Ok(config) = read_to_string(config) { let config = match toml::from_str(&config) { Ok(config) => config, Err(err) => return Err(Error::config(&config, err)), @@ -50,7 +50,7 @@ impl ConfigAndUsers { Config::default() }; - let users: Users = if let Ok(users) = read_to_string("users.toml") { + let users: Users = if let Ok(users) = read_to_string(users) { let users = toml::from_str(&users)?; info!("Loaded users.toml"); users diff --git a/pgdog/src/main.rs b/pgdog/src/main.rs index 269deb9..8e82104 100644 --- a/pgdog/src/main.rs +++ b/pgdog/src/main.rs @@ -7,6 +7,8 @@ use tokio::runtime::Builder; use tracing::info; use tracing_subscriber::{fmt, prelude::*, EnvFilter}; +use std::io::IsTerminal; + pub mod admin; pub mod auth; pub mod backend; @@ -20,16 +22,20 @@ pub mod state; pub mod stats; fn main() -> Result<(), Box> { + let args = cli::Cli::parse(); + + let format = fmt::layer() + .with_ansi(std::io::stderr().is_terminal()) + .with_file(false); + tracing_subscriber::registry() - .with(fmt::layer()) + .with(format) .with(EnvFilter::from_default_env()) .init(); - let args = cli::Cli::parse(); - info!("🐕 pgDog {}", env!("CARGO_PKG_VERSION")); - let config = config::load()?; + let config = config::load(&args.config, &args.users)?; plugin::load_from_config()?;