rullst_orm/
replica.rs

1use crate::RullstPool;
2use std::sync::atomic::{AtomicUsize, Ordering};
3
4/// Multi-Database Read Replica Load Balancer.
5/// Transparently load-balances read queries across secondary replicas while sending write queries to the primary database.
6pub struct ReplicaPool {
7    primary: RullstPool,
8    replicas: Vec<RullstPool>,
9    current_replica: AtomicUsize,
10}
11
12impl ReplicaPool {
13    /// Creates a new ReplicaPool with a primary pool and a list of read replica pools.
14    pub fn new(primary: RullstPool, replicas: Vec<RullstPool>) -> Self {
15        Self {
16            primary,
17            replicas,
18            current_replica: AtomicUsize::new(0),
19        }
20    }
21
22    /// Creates a ReplicaPool with only a primary pool (no read replicas).
23    pub fn primary_only(primary: RullstPool) -> Self {
24        Self::new(primary, Vec::new())
25    }
26
27    /// Returns a pool reference for read queries using Round-Robin balancing across replicas.
28    /// Falls back to the primary pool if no replicas are available.
29    pub fn read(&self) -> &RullstPool {
30        if self.replicas.is_empty() {
31            &self.primary
32        } else {
33            let idx = self.current_replica.fetch_add(1, Ordering::Relaxed) % self.replicas.len();
34            &self.replicas[idx]
35        }
36    }
37
38    /// Returns a pool reference to the primary database for write/mutation operations.
39    pub fn write(&self) -> &RullstPool {
40        &self.primary
41    }
42
43    /// Returns the number of active read replicas configured.
44    pub fn replica_count(&self) -> usize {
45        self.replicas.len()
46    }
47}
48
49#[cfg(test)]
50#[allow(unused_imports)]
51mod tests {
52    use super::*;
53
54    #[tokio::test]
55    async fn test_replica_pool_count_and_rotation() {
56        #[cfg(not(any(
57            feature = "strict-postgres",
58            feature = "strict-mysql",
59            feature = "strict-sqlite"
60        )))]
61        {
62            sqlx::any::install_default_drivers();
63        }
64
65        #[cfg(not(any(feature = "strict-postgres", feature = "strict-mysql")))]
66        {
67            if let (Ok(pool), Ok(replica1), Ok(replica2)) = (
68                RullstPool::connect("sqlite::memory:").await,
69                RullstPool::connect("sqlite::memory:").await,
70                RullstPool::connect("sqlite::memory:").await,
71            ) {
72                let replica_pool = ReplicaPool::new(pool.clone(), vec![replica1, replica2]);
73                assert_eq!(replica_pool.replica_count(), 2);
74
75                // Test read round-robin
76                let _r1 = replica_pool.read();
77                let _r2 = replica_pool.read();
78                let _r3 = replica_pool.read();
79
80                let primary_only = ReplicaPool::primary_only(pool);
81                assert_eq!(primary_only.replica_count(), 0);
82                let _p_read = primary_only.read();
83            }
84        }
85    }
86}
rullst_orm/replica.rs

rullst_orm/
replica.rs
