use crate::backend::{self, Backend, Connector};
use async_trait::async_trait;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
pub struct SlowConnector {
delay_ms: AtomicU64,
panic_on_access: AtomicBool,
}
impl SlowConnector {
pub fn new() -> Self {
Self {
delay_ms: AtomicU64::new(1),
panic_on_access: AtomicBool::new(false),
}
}
pub fn stall(&self) {
self.delay_ms.store(9999999, Ordering::SeqCst);
}
pub fn panic_on_access(&self) {
self.panic_on_access.store(true, Ordering::SeqCst);
}
async fn react_to_connection_operation(&self) {
if self.panic_on_access.load(Ordering::SeqCst) {
panic!("Should not be making new requests through this connector!");
}
let delay_ms = self.delay_ms.load(Ordering::SeqCst);
tokio::time::sleep(tokio::time::Duration::from_millis(delay_ms)).await;
}
}
#[async_trait]
impl Connector for SlowConnector {
type Connection = ();
async fn connect(&self, _backend: &Backend) -> Result<Self::Connection, backend::Error> {
self.react_to_connection_operation().await;
Ok(())
}
async fn is_valid(&self, _: &mut Self::Connection) -> Result<(), backend::Error> {
self.react_to_connection_operation().await;
Ok(())
}
async fn on_acquire(&self, _: &mut Self::Connection) -> Result<(), backend::Error> {
self.react_to_connection_operation().await;
Ok(())
}
async fn on_recycle(&self, _: &mut Self::Connection) -> Result<(), backend::Error> {
self.react_to_connection_operation().await;
Ok(())
}
}
pub struct FaultyConnector {
fail_on_access: AtomicBool,
}
impl FaultyConnector {
pub fn new() -> Self {
Self {
fail_on_access: AtomicBool::new(false),
}
}
pub fn start_failing(&self) {
self.fail_on_access.store(true, Ordering::SeqCst);
}
async fn react_to_connection_operation(&self) -> Result<(), backend::Error> {
if self.fail_on_access.load(Ordering::SeqCst) {
return Err(backend::Error::Other(anyhow::anyhow!(
"Connector failing intentionally"
)));
}
Ok(())
}
}
#[async_trait]
impl Connector for FaultyConnector {
type Connection = ();
async fn connect(&self, _backend: &Backend) -> Result<Self::Connection, backend::Error> {
self.react_to_connection_operation().await?;
Ok(())
}
async fn is_valid(&self, _: &mut Self::Connection) -> Result<(), backend::Error> {
self.react_to_connection_operation().await?;
Ok(())
}
async fn on_acquire(&self, _: &mut Self::Connection) -> Result<(), backend::Error> {
self.react_to_connection_operation().await?;
Ok(())
}
async fn on_recycle(&self, _: &mut Self::Connection) -> Result<(), backend::Error> {
self.react_to_connection_operation().await?;
Ok(())
}
}