diff --git a/src/pool.rs b/src/pool.rs index 07b7a40b..98c128af 100644 --- a/src/pool.rs +++ b/src/pool.rs @@ -12,7 +12,7 @@ use std::collections::{HashMap, HashSet}; use std::sync::Arc; use std::time::Instant; -use crate::config::{get_config, Address, PoolMode, Role, User}; +use crate::config::{get_config, Address, General, PoolMode, Role, User}; use crate::errors::Error; use crate::server::Server; @@ -82,6 +82,12 @@ pub struct PoolSettings { // Sharding key pub automatic_sharding_key: Option, + + // Health check timeout + pub healthcheck_timeout: u64, + + // Health check delay + pub healthcheck_delay: u64, } impl Default for PoolSettings { @@ -95,6 +101,8 @@ impl Default for PoolSettings { primary_reads_enabled: true, sharding_function: ShardingFunction::PgBigintHash, automatic_sharding_key: None, + healthcheck_delay: General::default_healthcheck_delay(), + healthcheck_timeout: General::default_healthcheck_timeout(), } } } @@ -256,6 +264,8 @@ impl ConnectionPool { primary_reads_enabled: pool_config.primary_reads_enabled, sharding_function: pool_config.sharding_function, automatic_sharding_key: pool_config.automatic_sharding_key.clone(), + healthcheck_delay: config.general.healthcheck_delay, + healthcheck_timeout: config.general.healthcheck_timeout, }, }; @@ -343,9 +353,6 @@ impl ConnectionPool { // Random load balancing candidates.shuffle(&mut thread_rng()); - let healthcheck_timeout = get_config().general.healthcheck_timeout; - let healthcheck_delay = get_config().general.healthcheck_delay as u128; - while !candidates.is_empty() { // Get the next candidate let address = match candidates.pop() { @@ -380,8 +387,8 @@ impl ConnectionPool { let server = &mut *conn; // Will return error if timestamp is greater than current system time, which it should never be set to - let require_healthcheck = - server.last_activity().elapsed().unwrap().as_millis() > healthcheck_delay; + let require_healthcheck = server.last_activity().elapsed().unwrap().as_millis() + > self.settings.healthcheck_delay as u128; // Do not issue a health check unless it's been a little while // since we last checked the server is ok. @@ -398,7 +405,7 @@ impl ConnectionPool { self.stats.server_tested(server.server_id()); match tokio::time::timeout( - tokio::time::Duration::from_millis(healthcheck_timeout), + tokio::time::Duration::from_millis(self.settings.healthcheck_timeout), server.query(";"), // Cheap query as it skips the query planner ) .await diff --git a/src/query_router.rs b/src/query_router.rs index 552c358c..50905716 100644 --- a/src/query_router.rs +++ b/src/query_router.rs @@ -775,6 +775,8 @@ mod test { primary_reads_enabled: false, sharding_function: ShardingFunction::PgBigintHash, automatic_sharding_key: Some(String::from("id")), + healthcheck_delay: PoolSettings::default().healthcheck_delay, + healthcheck_timeout: PoolSettings::default().healthcheck_timeout, }; let mut qr = QueryRouter::new(); assert_eq!(qr.active_role, None);