diff --git a/martin/src/pg/config.rs b/martin/src/pg/config.rs index b5457d886..1a46bb90c 100644 --- a/martin/src/pg/config.rs +++ b/martin/src/pg/config.rs @@ -1,8 +1,8 @@ +use std::time::Duration; + use futures::future::try_join; -use futures::pin_mut; use log::warn; use serde::{Deserialize, Serialize}; -use std::time::Duration; use tilejson::TileJSON; use crate::config::{copy_unrecognized_config, UnrecognizedValues}; @@ -11,7 +11,7 @@ use crate::pg::config_table::TableInfoSources; use crate::pg::configurator::PgBuilder; use crate::pg::Result; use crate::source::Sources; -use crate::utils::{sorted_opt_map, BoolOrObject, IdResolver, OneOrMany}; +use crate::utils::{on_slow, sorted_opt_map, BoolOrObject, IdResolver, OneOrMany}; pub trait PgInfo { fn format_id(&self) -> String; @@ -113,31 +113,21 @@ impl PgConfig { pub async fn resolve(&mut self, id_resolver: IdResolver) -> crate::Result { let pg = PgBuilder::new(self, id_resolver).await?; + let inst_tables = on_slow(pg.instantiate_tables(), Duration::from_secs(5), || { + if pg.disable_bounds() { + warn!("Discovering tables in PostgreSQL database '{}' is taking too long. Bounds calculation is already disabled. You may need to tune your database.", pg.get_id()); + } else { + warn!("Discovering tables in PostgreSQL database '{}' is taking too long. Make sure your table geo columns have a GIS index, or use --disable-bounds CLI/config to skip bbox calculation.", pg.get_id()); + } + }); let ((mut tables, tbl_info), (funcs, func_info)) = - try_join(Self::instantiate_tables(&pg), pg.instantiate_functions()).await?; + try_join(inst_tables, pg.instantiate_functions()).await?; self.tables = Some(tbl_info); self.functions = Some(func_info); tables.extend(funcs); Ok(tables) } - - async fn instantiate_tables(pg: &PgBuilder) -> Result<(Sources, TableInfoSources)> { - let instantiate_tables = pg.instantiate_tables(); - pin_mut!(instantiate_tables); - - let timeout_future = tokio::time::sleep(Duration::from_secs(5)); - - tokio::select! { - result = &mut instantiate_tables => { - result - } - () = timeout_future => { - warn!("PostgreSQL table discovery is taking too long. Make sure your table geo columns have a GIS index, or use --disabling-bounds."); - instantiate_tables.await - } - } - } } #[cfg(test)] diff --git a/martin/src/pg/configurator.rs b/martin/src/pg/configurator.rs index d5fbf77d0..5acdeb65b 100755 --- a/martin/src/pg/configurator.rs +++ b/martin/src/pg/configurator.rs @@ -63,6 +63,14 @@ impl PgBuilder { }) } + pub fn disable_bounds(&self) -> bool { + self.disable_bounds + } + + pub fn get_id(&self) -> &str { + self.pool.get_id() + } + // FIXME: this function has gotten too long due to the new formatting rules, need to be refactored #[allow(clippy::too_many_lines)] pub async fn instantiate_tables(&self) -> Result<(Sources, TableInfoSources)> { diff --git a/martin/src/utils/utilities.rs b/martin/src/utils/utilities.rs index c7c67d2ed..6635990bb 100644 --- a/martin/src/utils/utilities.rs +++ b/martin/src/utils/utilities.rs @@ -1,9 +1,13 @@ use std::collections::{BTreeMap, HashMap}; +use std::future::Future; use std::io::{Read as _, Write as _}; +use std::time::Duration; use flate2::read::GzDecoder; use flate2::write::GzEncoder; +use futures::pin_mut; use serde::{Deserialize, Serialize, Serializer}; +use tokio::time::timeout; #[must_use] pub fn is_valid_zoom(zoom: u8, minzoom: Option, maxzoom: Option) -> bool { @@ -58,3 +62,17 @@ pub fn encode_brotli(data: &[u8]) -> Result, std::io::Error> { encoder.write_all(data)?; Ok(encoder.into_inner()) } + +pub async fn on_slow( + future: impl Future, + duration: Duration, + fn_on_slow: S, +) -> T { + pin_mut!(future); + if let Ok(result) = timeout(duration, &mut future).await { + result + } else { + fn_on_slow(); + future.await + } +}