From 321a969a39c765cfe936ab9b3d4cdb982d02659e Mon Sep 17 00:00:00 2001
From: Pat Sier <pjsier@users.noreply.github.com>
Date: Sun, 8 Oct 2023 16:13:00 -0400
Subject: [PATCH] Log a warning if PG startup takes too long (#924)

Adds a warning message using the suggested implementation in #810.
Thanks for the recommended approach in the ticket! This one seemed
straightforward enough that local tests would do, which I ran by adding
a delay to `instantiate_tables`.

---------

Co-authored-by: Yuri Astrakhan <YuriAstrakhan@gmail.com>
---
 martin/src/pg/config.rs       | 14 ++++++++++++--
 martin/src/pg/configurator.rs |  8 ++++++++
 martin/src/utils/utilities.rs | 18 ++++++++++++++++++
 3 files changed, 38 insertions(+), 2 deletions(-)

diff --git a/martin/src/pg/config.rs b/martin/src/pg/config.rs
index c84aa01b3..1a46bb90c 100644
--- a/martin/src/pg/config.rs
+++ b/martin/src/pg/config.rs
@@ -1,4 +1,7 @@
+use std::time::Duration;
+
 use futures::future::try_join;
+use log::warn;
 use serde::{Deserialize, Serialize};
 use tilejson::TileJSON;
 
@@ -8,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;
@@ -110,8 +113,15 @@ 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(pg.instantiate_tables(), pg.instantiate_functions()).await?;
+            try_join(inst_tables, pg.instantiate_functions()).await?;
 
         self.tables = Some(tbl_info);
         self.functions = Some(func_info);
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<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
+    }
+}