remove prefered order feature

MarcosNicolau · MarcosNicolau · commit 4a67e5872091 · 2026-01-05T19:01:27.000-03:00
diff --git a/aggregation_mode/db/src/orchestrator.rs b/aggregation_mode/db/src/orchestrator.rs
@@ -1,11 +1,4 @@
-use std::{
-    future::Future,
-    sync::{
-        atomic::{AtomicBool, Ordering},
-        Arc,
-    },
-    time::Duration,
-};
+use std::{future::Future, sync::Arc, time::Duration};
 
 use sqlx::{postgres::PgPoolOptions, Pool, Postgres};
 
@@ -21,8 +14,6 @@ enum Operation {
 #[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.
@@ -77,11 +68,7 @@ impl DbOrchestrator {
             .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),
-                }))
+                Ok(Arc::new(DbNode { pool }))
             })
             .collect::<Result<Vec<_>, sqlx::Error>>()
             .map_err(DbOrchestratorError::Sqlx)?;
@@ -92,23 +79,7 @@ impl DbOrchestrator {
         })
     }
 
-    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>
+    pub async fn query<T, Q, Fut>(&self, query_fn: Q) -> Result<T, sqlx::Error>
     where
         Q: Fn(Pool<Postgres>) -> Fut,
         Fut: Future<Output = Result<T, sqlx::Error>>,
@@ -117,7 +88,7 @@ impl DbOrchestrator {
         let mut delay = Duration::from_millis(self.retry_config.min_delay_millis);
 
         loop {
-            match self.execute_once(&query_fn, operation).await {
+            match self.execute_once(&query_fn).await {
                 Ok(value) => return Ok(value),
                 Err(RetryError::Permanent(err)) => return Err(err),
                 Err(RetryError::Transient(err)) => {
@@ -134,38 +105,23 @@ impl DbOrchestrator {
         }
     }
 
-    async fn execute_once<T, Q, Fut>(
-        &self,
-        query_fn: &Q,
-        operation: Operation,
-    ) -> Result<T, RetryError<sqlx::Error>>
+    async fn execute_once<T, Q, Fut>(&self, query_fn: &Q) -> 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];
+        for (idx, node) in self.nodes.iter().enumerate() {
             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));
@@ -179,27 +135,6 @@ impl DbOrchestrator {
         ))
     }
 
-    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,
diff --git a/aggregation_mode/gateway/src/db.rs b/aggregation_mode/gateway/src/db.rs
@@ -60,7 +60,7 @@ impl Db {
 
     pub async fn count_tasks_by_address(&self, address: &str) -> Result<i64, sqlx::Error> {
         self.orchestrator
-            .read(async |pool| {
+            .query(async |pool| {
                 let (count,) =
                     sqlx::query_as::<_, (i64,)>("SELECT COUNT(*) FROM tasks WHERE address = $1")
                         .bind(address.to_lowercase())
@@ -77,7 +77,7 @@ impl Db {
         task_id: Uuid,
     ) -> Result<Option<Vec<u8>>, sqlx::Error> {
         self.orchestrator
-            .read(async |pool| {
+            .query(async |pool| {
                 sqlx::query_scalar::<_, Option<Vec<u8>>>(
                     "SELECT merkle_path FROM tasks WHERE task_id = $1",
                 )
@@ -95,7 +95,7 @@ impl Db {
         nonce: i64,
     ) -> Result<Vec<Receipt>, sqlx::Error> {
         self.orchestrator
-            .read(async |pool| {
+            .query(async |pool| {
                 sqlx::query_as::<_, Receipt>(
                     "SELECT status,merkle_path,nonce,address FROM tasks
                     WHERE address = $1
@@ -116,7 +116,7 @@ impl Db {
         limit: i64,
     ) -> Result<Vec<Receipt>, sqlx::Error> {
         self.orchestrator
-            .read(async |pool| {
+            .query(async |pool| {
                 sqlx::query_as::<_, Receipt>(
                     "SELECT status,merkle_path,nonce,address FROM tasks
                     WHERE address = $1
@@ -133,7 +133,7 @@ impl Db {
 
     pub async fn get_daily_tasks_by_address(&self, address: &str) -> Result<i64, sqlx::Error> {
         self.orchestrator
-            .read(async |pool| {
+            .query(async |pool| {
                 sqlx::query_scalar::<_, i64>(
                     "SELECT COUNT(*)
                 FROM tasks
@@ -157,7 +157,7 @@ impl Db {
         nonce: i64,
     ) -> Result<Uuid, sqlx::Error> {
         self.orchestrator
-            .write(async |pool| {
+            .query(async |pool| {
                 sqlx::query_scalar::<_, Uuid>(
                     "INSERT INTO tasks (
                         address,
@@ -187,7 +187,7 @@ impl Db {
         epoch: BigDecimal,
     ) -> Result<bool, sqlx::Error> {
         self.orchestrator
-            .read(async |pool| {
+            .query(async |pool| {
                 sqlx::query_scalar::<_, bool>(
                     "SELECT EXISTS (
                     SELECT 1 FROM payment_events
diff --git a/aggregation_mode/payments_poller/src/db.rs b/aggregation_mode/payments_poller/src/db.rs
@@ -62,7 +62,7 @@ impl Db {
         tx_hash: &str,
     ) -> Result<(), sqlx::Error> {
         self.orchestrator
-            .write(async |pool| {
+            .query(async |pool| {
                 sqlx::query(
                     "INSERT INTO payment_events (address, started_at, amount, valid_until, tx_hash)
                 VALUES ($1, $2, $3, $4, $5)
diff --git a/aggregation_mode/proof_aggregator/src/backend/db.rs b/aggregation_mode/proof_aggregator/src/backend/db.rs
@@ -60,7 +60,7 @@ impl Db {
         limit: i64,
     ) -> Result<Vec<Task>, DbError> {
         self.orchestrator
-            .write(async |pool| {
+            .query(async |pool| {
                 sqlx::query_as::<_, Task>(
                     "WITH selected AS (
                         SELECT task_id
@@ -91,7 +91,7 @@ impl Db {
         let updates_ref = &updates;
 
         self.orchestrator
-            .write(|pool| {
+            .query(|pool| {
                 let updates = updates_ref;
                 async move {
                     let mut tx = pool.begin().await?;
