From 953f23f674f1cca64f7a77fdffd93119e219f423 Mon Sep 17 00:00:00 2001
From: Theo <theo.tanadi122@gmail.com>
Date: Mon, 21 Oct 2024 14:41:48 +0700
Subject: [PATCH] Add retry mechanism to connect to PostgreSQL before starting
 maplibre (#1539)

---
 Cargo.lock              | 12 ++++++++++++
 martin/Cargo.toml       |  1 +
 martin/src/pg/config.rs | 17 +++++++++++++++--
 3 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 4c796d47..7964eb2e 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2544,6 +2544,7 @@ dependencies = [
  "tilejson",
  "tokio",
  "tokio-postgres-rustls",
+ "tokio-retry",
  "url",
 ]
 
@@ -4948,6 +4949,17 @@ dependencies = [
  "x509-cert",
 ]
 
+[[package]]
+name = "tokio-retry"
+version = "0.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7f57eb36ecbe0fc510036adff84824dd3c24bb781e21bfa67b69d556aa85214f"
+dependencies = [
+ "pin-project",
+ "rand",
+ "tokio",
+]
+
 [[package]]
 name = "tokio-rustls"
 version = "0.26.0"
diff --git a/martin/Cargo.toml b/martin/Cargo.toml
index 28dc1c4d..169c3a30 100644
--- a/martin/Cargo.toml
+++ b/martin/Cargo.toml
@@ -113,6 +113,7 @@ tilejson.workspace = true
 tokio = { workspace = true, features = ["io-std"] }
 tokio-postgres-rustls = { workspace = true, optional = true }
 url.workspace = true
+tokio-retry = "0.3.0"
 
 [build-dependencies]
 static-files = { workspace = true, optional = true }
diff --git a/martin/src/pg/config.rs b/martin/src/pg/config.rs
index 39ad64d9..ef2d2eed 100644
--- a/martin/src/pg/config.rs
+++ b/martin/src/pg/config.rs
@@ -15,7 +15,9 @@ use crate::pg::utils::on_slow;
 use crate::pg::PgResult;
 use crate::source::TileInfoSources;
 use crate::utils::{IdResolver, OptBoolObj, OptOneMany};
-use crate::MartinResult;
+use crate::{MartinError, MartinResult};
+use tokio_retry::strategy::{jitter, FixedInterval};
+use tokio_retry::Retry;
 
 pub trait PgInfo {
     fn format_id(&self) -> String;
@@ -114,7 +116,18 @@ impl PgConfig {
     }
 
     pub async fn resolve(&mut self, id_resolver: IdResolver) -> MartinResult<TileInfoSources> {
-        let pg = PgBuilder::new(self, id_resolver).await?;
+        // Retry strategy: Fixed 5 seconds interval backoff with jitter (random variation)
+        let retry_strategy = FixedInterval::from_millis(5000)
+            .map(jitter) // Add random jitter to avoid "thundering herd" problem
+            .take(3); // Retry up to 3 times
+
+        // Create PgBuilder using retry_strategy
+        let pg = Retry::spawn(retry_strategy, || async {
+            PgBuilder::new(self, id_resolver.clone()).await
+        })
+        .await
+        .map_err(MartinError::PostgresError)?;
+
         let inst_tables = on_slow(
             pg.instantiate_tables(),
             // warn only if default bounds timeout has already passed