Merge pull request #1 from nyurik/810-fix

Cleaner approach to timeout reporting

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<Sources> {
         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)]

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)> {

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<u8>, maxzoom: Option<u8>) -> bool {
@@ -58,3 +62,17 @@ pub fn encode_brotli(data: &[u8]) -> Result<Vec<u8>, std::io::Error> {
     encoder.write_all(data)?;
     Ok(encoder.into_inner())
 }
+
+pub async fn on_slow<T, S: FnOnce()>(
+    future: impl Future<Output = T>,
+    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
+    }
+}