main
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use sea_orm::{DatabaseConnection, DbErr};
#[derive(Clone)]
pub struct DbConn {
write: DatabaseConnection,
read: Arc<[DatabaseConnection]>,
read_index: Arc<AtomicUsize>,
pub(crate) replica_failure_logged: Arc<AtomicBool>,
}
impl DbConn {
pub fn new(write: DatabaseConnection, read: Vec<DatabaseConnection>) -> Self {
Self {
write,
read: read.into(),
read_index: Arc::new(AtomicUsize::new(0)),
replica_failure_logged: Arc::new(AtomicBool::new(false)),
}
}
pub fn write(&self) -> &DatabaseConnection {
&self.write
}
pub fn read(&self) -> &Self {
self
}
pub fn replicas(&self) -> &[DatabaseConnection] {
&self.read
}
pub(crate) fn replica(&self) -> Option<&DatabaseConnection> {
if self.read.is_empty() {
return None;
}
let index = self.read_index.fetch_add(1, Ordering::Relaxed) % self.read.len();
Some(&self.read[index])
}
pub async fn close(&self) -> Result<(), DbErr> {
let mut result = self.write.close_by_ref().await;
for connection in self.read.iter() {
let replica_result = connection.close_by_ref().await;
if result.is_ok() {
result = replica_result;
}
}
result
}
}