feat(aggregation-mode): support multiple dbs

aggregation_mode/db/Cargo.toml

@@ -4,10 +4,11 @@ version = "0.1.0"
 edition = "2021"
 
 [dependencies]
+serde =  { workspace = true }
 tokio = { version = "1"}
-# TODO: enable tls
 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/src/lib.rs

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

aggregation_mode/db/src/orchestrator.rs

@@ -0,0 +1,228 @@
+use std::{
+    future::Future,
+    sync::{
+        atomic::{AtomicBool, Ordering},
+        Arc,
+    },
+    time::Duration,
+};
+
+use sqlx::{postgres::PgPoolOptions, Pool, Postgres};
+
+use crate::retry::{RetryConfig, RetryError};
+
+#[derive(Debug, Clone, Copy)]
+enum Operation {
+    Read,
+    Write,
+}
+
+/// A single DB node: connection pool plus shared health flags (used to prioritize nodes).
+
+#[derive(Debug)]
+struct DbNode {
+    pool: Pool<Postgres>,
+    last_read_failed: AtomicBool,
+    last_write_failed: AtomicBool,
+}
+
+/// Database orchestrator for running reads/writes across multiple PostgreSQL nodes with retry/backoff.
+///
+/// `DbOrchestartor` holds a list of database nodes (connection pools) and will:
+/// - try nodes in a preferred order (healthy nodes first, then recently-failed nodes),
+/// - mark nodes as failed on connection-type errors,
+/// - retry transient failures with exponential backoff based on `retry_config`,
+///
+/// ## Thread-safe `Clone`
+/// This type is cheap and thread-safe to clone:
+/// - `nodes` is `Vec<Arc<DbNode>>`, so cloning only increments `Arc` ref-counts and shares the same pools/nodes,
+/// - `sqlx::Pool<Postgres>` is internally reference-counted and designed to be cloned and used concurrently,
+/// - the node health flags are `AtomicBool`, so updates are safe from multiple threads/tasks.
+///
+/// Clones share health state (the atomics) and the underlying pools, so all clones observe and influence
+/// the same “preferred node” ordering decisions.
+#[derive(Debug, Clone)]
+pub struct DbOrchestartor {
+    nodes: Vec<Arc<DbNode>>,
+    retry_config: RetryConfig,
+}
+
+#[derive(Debug)]
+pub enum DbOrchestartorError {
+    InvalidNumberOfConnectionUrls,
+    Sqlx(sqlx::Error),
+}
+
+impl std::fmt::Display for DbOrchestartorError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Self::InvalidNumberOfConnectionUrls => {
+                write!(f, "invalid number of connection URLs")
+            }
+            Self::Sqlx(e) => write!(f, "{e}"),
+        }
+    }
+}
+
+impl DbOrchestartor {
+    pub fn try_new(
+        connection_urls: &[String],
+        retry_config: RetryConfig,
+    ) -> Result<Self, DbOrchestartorError> {
+        if connection_urls.is_empty() {
+            return Err(DbOrchestartorError::InvalidNumberOfConnectionUrls);
+        }
+
+        let nodes = connection_urls
+            .iter()
+            .map(|url| {
+                let pool = PgPoolOptions::new().max_connections(5).connect_lazy(url)?;
+
+                Ok(Arc::new(DbNode {
+                    pool,
+                    last_read_failed: AtomicBool::new(false),
+                    last_write_failed: AtomicBool::new(false),
+                }))
+            })
+            .collect::<Result<Vec<_>, sqlx::Error>>()
+            .map_err(DbOrchestartorError::Sqlx)?;
+
+        Ok(Self {
+            nodes,
+            retry_config,
+        })
+    }
+
+    pub async fn write<T, Q, Fut>(&self, query: Q) -> Result<T, sqlx::Error>
+    where
+        Q: Fn(Pool<Postgres>) -> Fut,
+        Fut: Future<Output = Result<T, sqlx::Error>>,
+    {
+        self.query::<T, Q, Fut>(query, Operation::Write).await
+    }
+
+    pub async fn read<T, Q, Fut>(&self, query: Q) -> Result<T, sqlx::Error>
+    where
+        Q: Fn(Pool<Postgres>) -> Fut,
+        Fut: Future<Output = Result<T, sqlx::Error>>,
+    {
+        self.query::<T, Q, Fut>(query, Operation::Read).await
+    }
+
+    async fn query<T, Q, Fut>(&self, query_fn: Q, operation: Operation) -> Result<T, sqlx::Error>
+    where
+        Q: Fn(Pool<Postgres>) -> Fut,
+        Fut: Future<Output = Result<T, sqlx::Error>>,
+    {
+        let mut attempts = 0;
+        let mut delay = Duration::from_millis(self.retry_config.min_delay_millis);
+
+        loop {
+            match self.execute_once(&query_fn, operation).await {
+                Ok(value) => return Ok(value),
+                Err(RetryError::Permanent(err)) => return Err(err),
+                Err(RetryError::Transient(err)) => {
+                    if attempts >= self.retry_config.max_delay_seconds {
+                        return Err(err);
+                    }
+
+                    tracing::warn!(attempt = attempts, delay_milis = delay.as_millis(), error = ?err, "retrying after backoff");
+                    tokio::time::sleep(delay).await;
+                    delay = self.backoff_delay(delay);
+                    attempts += 1;
+                }
+            }
+        }
+    }
+
+    fn backoff_delay(&self, current: Duration) -> Duration {
+        let max = Duration::from_secs(self.retry_config.max_delay_seconds);
+        let scaled_secs = current.as_secs_f64() * f64::from(self.retry_config.factor);
+        let scaled = Duration::from_secs_f64(scaled_secs);
+        if scaled > max {
+            max
+        } else {
+            scaled
+        }
+    }
+
+    async fn execute_once<T, Q, Fut>(
+        &self,
+        query_fn: &Q,
+        operation: Operation,
+    ) -> Result<T, RetryError<sqlx::Error>>
+    where
+        Q: Fn(Pool<Postgres>) -> Fut,
+        Fut: Future<Output = Result<T, sqlx::Error>>,
+    {
+        let mut last_error = None;
+
+        for idx in self.preferred_order(operation) {
+            let node = &self.nodes[idx];
+            let pool = node.pool.clone();
+
+            match query_fn(pool).await {
+                Ok(res) => {
+                    match operation {
+                        Operation::Read => node.last_read_failed.store(false, Ordering::Relaxed),
+                        Operation::Write => node.last_write_failed.store(false, Ordering::Relaxed),
+                    };
+                    return Ok(res);
+                }
+                Err(err) => {
+                    if Self::is_connection_error(&err) {
+                        tracing::warn!(node_index = idx, error = ?err, "database query failed");
+                        match operation {
+                            Operation::Read => node.last_read_failed.store(true, Ordering::Relaxed),
+                            Operation::Write => {
+                                node.last_write_failed.store(true, Ordering::Relaxed)
+                            }
+                        };
+                        last_error = Some(err);
+                    } else {
+                        return Err(RetryError::Permanent(err));
+                    }
+                }
+            };
+        }
+
+        Err(RetryError::Transient(
+            last_error.expect("write_op attempted without database nodes"),
+        ))
+    }
+
+    fn preferred_order(&self, operation: Operation) -> Vec<usize> {
+        let mut preferred = Vec::with_capacity(self.nodes.len());
+        let mut fallback = Vec::new();
+
+        for (idx, node) in self.nodes.iter().enumerate() {
+            let failed = match operation {
+                Operation::Read => node.last_read_failed.load(Ordering::Relaxed),
+                Operation::Write => node.last_write_failed.load(Ordering::Relaxed),
+            };
+
+            if failed {
+                fallback.push(idx);
+            } else {
+                preferred.push(idx);
+            }
+        }
+
+        preferred.extend(fallback);
+        preferred
+    }
+
+    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
+                | sqlx::Error::Database(_)
+        )
+    }
+}

aggregation_mode/db/src/retry.rs

@@ -0,0 +1,28 @@
+#[derive(Debug)]
+pub(super) 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: std::fmt::Display> std::error::Error for RetryError<E> where E: std::fmt::Debug {}
+
+#[derive(Debug, Clone)]
+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,
+}

aggregation_mode/db/src/types.rs

@@ -4,7 +4,7 @@ use sqlx::{
     Type,
 };
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Type)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Type, serde::Serialize)]
 #[sqlx(type_name = "task_status", rename_all = "lowercase")]
 pub enum TaskStatus {
     Pending,

aggregation_mode/gateway/Cargo.toml

@@ -10,6 +10,7 @@ serde_yaml = { workspace = true }
 agg_mode_sdk = { path = "../sdk"}
 aligned-sdk = { workspace = true }
 sp1-sdk = { workspace = true }
+db = { workspace = true }
 tracing = { version = "0.1", features = ["log"] }
 tracing-subscriber = { version = "0.3.0", features = ["env-filter"] }
 bincode = "1.3.3"

aggregation_mode/gateway/src/config.rs

@@ -5,7 +5,7 @@ use serde::{Deserialize, Serialize};
 #[derive(Clone, Debug, Deserialize, Serialize)]
 pub struct Config {
     pub port: u16,
-    pub db_connection_url: String,
+    pub db_connection_urls: Vec<String>,
     pub network: String,
     pub max_daily_proofs_per_user: i64,
 }