feat: db orchestrator design

Author: MarcosNicolau

aggregation_mode/Cargo.lock

@@ -5,9 +5,8 @@ edition = "2021"
 
 [dependencies]
 tokio = { version = "1"}
-# TODO: enable tls
 sqlx = { version = "0.8", features = [ "runtime-tokio", "postgres", "migrate" ] }
-
+backon = "1.2.0"
 
 [[bin]]
 name = "migrate"

aggregation_mode/db/Cargo.toml

@@ -1 +1,3 @@
+pub mod orchestrator;
+mod retry;
 pub mod types;

aggregation_mode/db/src/lib.rs

@@ -0,0 +1,86 @@
+use std::{future::Future, time::Duration};
+
+use backon::{ExponentialBuilder, Retryable};
+use sqlx::{postgres::PgPoolOptions, Pool, Postgres};
+
+use crate::retry::RetryError;
+
+#[derive(Clone, Debug)]
+struct DbNode {
+    pool: Pool<Postgres>,
+}
+
+pub struct DbOrchestartor {
+    nodes: Vec<DbNode>,
+}
+
+impl DbOrchestartor {
+    pub fn try_new(connection_urls: Vec<String>) -> Result<Self, sqlx::Error> {
+        // TODO: validate at least one connection url
+        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>>()?;
+
+        Ok(Self { nodes })
+    }
+
+    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))
+    }
+
+    pub async fn query<T, E, Q, Fut>(&self, query_fn: Q) -> Result<T, sqlx::Error>
+    where
+        Q: Fn(Pool<Postgres>) -> Fut,
+        Fut: Future<Output = Result<T, sqlx::Error>>,
+    {
+        let func = async || {
+            let mut last_error = None;
+
+            for idx in 0..self.nodes.len() {
+                let pool = self.nodes[idx].pool.clone();
+                match query_fn(pool).await {
+                    Ok(res) => return Ok(res),
+                    Err(err) => {
+                        if Self::is_connection_error(&err) {
+                            last_error = Some(err);
+                        } else {
+                            return Err(RetryError::Permanent(err));
+                        }
+                    }
+                };
+            }
+
+            Err(RetryError::Transient(
+                last_error.expect("write_op attempted without database nodes"),
+            ))
+        };
+
+        func.retry(self.backoff_builder())
+            .sleep(tokio::time::sleep)
+            .when(|e| matches!(e, RetryError::Transient(_)))
+            .await
+            .map_err(|e| e.inner())
+    }
+
+    fn is_connection_error(error: &sqlx::Error) -> bool {
+        matches!(
+            error,
+            sqlx::Error::Io(_)
+                | sqlx::Error::Tls(_)
+                | sqlx::Error::Protocol(_)
+                | sqlx::Error::PoolTimedOut
+                | sqlx::Error::PoolClosed
+                | sqlx::Error::WorkerCrashed
+                | sqlx::Error::BeginFailed
+        )
+    }
+}

aggregation_mode/db/src/orchestrator.rs

@@ -0,0 +1,67 @@
+use backon::ExponentialBuilder;
+use backon::Retryable;
+use std::{future::Future, time::Duration};
+
+#[derive(Debug)]
+pub enum RetryError<E> {
+    Transient(E),
+    Permanent(E),
+}
+
+impl<E: std::fmt::Display> std::fmt::Display for RetryError<E> {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        match self {
+            RetryError::Transient(e) => write!(f, "{e}"),
+            RetryError::Permanent(e) => write!(f, "{e}"),
+        }
+    }
+}
+
+impl<E> RetryError<E> {
+    pub fn inner(self) -> E {
+        match self {
+            RetryError::Transient(e) => e,
+            RetryError::Permanent(e) => 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
+}