feat: receive config via params and log errors

Author: MarcosNicolau

aggregation_mode/Cargo.lock

@@ -7,6 +7,7 @@ edition = "2021"
 tokio = { version = "1"}
 sqlx = { version = "0.8", features = [ "runtime-tokio", "postgres", "migrate" ] }
 backon = "1.2.0"
+tracing = { version = "0.1", features = ["log"] }
 
 [[bin]]
 name = "migrate"

aggregation_mode/db/Cargo.toml

@@ -3,7 +3,7 @@ use std::{future::Future, time::Duration};
 use backon::{ExponentialBuilder, Retryable};
 use sqlx::{postgres::PgPoolOptions, Pool, Postgres};
 
-use crate::retry::RetryError;
+use crate::retry::{RetryConfig, RetryError};
 
 #[derive(Clone, Debug)]
 struct DbNode {
@@ -12,29 +12,45 @@ struct DbNode {
 
 pub struct DbOrchestartor {
     nodes: Vec<DbNode>,
+    retry_config: RetryConfig,
+}
+
+pub enum DbOrchestartorError {
+    InvalidNumberOfConnectionUrls,
+    Sqlx(sqlx::Error),
 }
 
 impl DbOrchestartor {
-    pub fn try_new(connection_urls: Vec<String>) -> Result<Self, sqlx::Error> {
-        // TODO: validate at least one connection url
+    pub fn try_new(
+        connection_urls: Vec<String>,
+        retry_config: RetryConfig,
+    ) -> Result<Self, DbOrchestartorError> {
+        if connection_urls.is_empty() {
+            return Err(DbOrchestartorError::InvalidNumberOfConnectionUrls);
+        }
+
         let nodes = connection_urls
             .into_iter()
             .map(|url| {
                 let pool = PgPoolOptions::new().max_connections(5).connect_lazy(&url)?;
 
                 Ok(DbNode { pool })
             })
-            .collect::<Result<Vec<_>, sqlx::Error>>()?;
+            .collect::<Result<Vec<_>, sqlx::Error>>()
+            .map_err(|e| DbOrchestartorError::Sqlx(e))?;
 
-        Ok(Self { nodes })
+        Ok(Self {
+            nodes,
+            retry_config,
+        })
     }
 
     fn backoff_builder(&self) -> ExponentialBuilder {
         ExponentialBuilder::default()
-            .with_min_delay(Duration::from_millis(0))
-            .with_max_times(0)
-            .with_factor(0.0)
-            .with_max_delay(Duration::from_secs(0))
+            .with_min_delay(Duration::from_millis(self.retry_config.min_delay_millis))
+            .with_max_times(self.retry_config.max_times)
+            .with_factor(self.retry_config.factor)
+            .with_max_delay(Duration::from_secs(self.retry_config.max_delay_seconds))
     }
 
     pub async fn query<T, E, Q, Fut>(&self, query_fn: Q) -> Result<T, sqlx::Error>
@@ -51,6 +67,7 @@ impl DbOrchestartor {
                     Ok(res) => return Ok(res),
                     Err(err) => {
                         if Self::is_connection_error(&err) {
+                            tracing::warn!(node_index = idx, error = ?err, "database query failed; retrying");
                             last_error = Some(err);
                         } else {
                             return Err(RetryError::Permanent(err));
@@ -81,6 +98,7 @@ impl DbOrchestartor {
                 | sqlx::Error::PoolClosed
                 | sqlx::Error::WorkerCrashed
                 | sqlx::Error::BeginFailed
+                | sqlx::Error::Database(_)
         )
     }
 }

aggregation_mode/db/src/orchestrator.rs

@@ -1,7 +1,3 @@
-use backon::ExponentialBuilder;
-use backon::Retryable;
-use std::{future::Future, time::Duration};
-
 #[derive(Debug)]
 pub enum RetryError<E> {
     Transient(E),
@@ -28,40 +24,14 @@ impl<E> RetryError<E> {
 
 impl<E: std::fmt::Display> std::error::Error for RetryError<E> where E: std::fmt::Debug {}
 
-/// Supports retries only on async functions. See: https://docs.rs/backon/latest/backon/#retry-an-async-function
-/// Runs with `jitter: false`.
-///
-/// # Parameters
-/// * `function` - The async function to retry
-/// * `min_delay_millis` - Initial delay before first retry attempt (in milliseconds)
-/// * `factor` - Exponential backoff multiplier for retry delays
-/// * `max_times` - Maximum number of retry attempts
-/// * `max_delay_seconds` - Maximum delay between retry attempts (in seconds)
-pub async fn retry_function<FutureFn, Fut, T, E, A>(
-    mut function: FutureFn,
-    min_delay_millis: u64,
-    factor: f32,
-    max_times: usize,
-    max_delay_seconds: u64,
-) -> Result<T, RetryError<E>>
-where
-    Fut: Future<Output = Result<T, RetryError<E>>>,
-    FutureFn: FnMut() -> Fut,
-{
-    let backoff = ExponentialBuilder::default()
-        .with_min_delay(Duration::from_millis(min_delay_millis))
-        .with_max_times(max_times)
-        .with_factor(factor)
-        .with_max_delay(Duration::from_secs(max_delay_seconds));
-
-    let func = async || match function().await {
-        Ok(res) => Ok(res),
-        Err(e) => Err(e),
-    };
-
-    function
-        .retry(backoff)
-        .sleep(tokio::time::sleep)
-        .when(|e| matches!(e, RetryError::Transient(_)))
-        .await
+#[derive(Debug)]
+pub struct RetryConfig {
+    /// * `min_delay_millis` - Initial delay before first retry attempt (in milliseconds)
+    pub min_delay_millis: u64,
+    /// * `factor` - Exponential backoff multiplier for retry delays
+    pub factor: f32,
+    /// * `max_times` - Maximum number of retry attempts
+    pub max_times: usize,
+    /// * `max_delay_seconds` - Maximum delay between retry attempts (in seconds)
+    pub max_delay_seconds: u64,
 }