use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use moka::future::Cache;
use tokio::sync::{OwnedSemaphorePermit, Semaphore};
use crate::error::EngineError;
use crate::models::{ImageSearchResponse, SearchResponse};
#[derive(Debug, Clone)]
pub enum CachedResponse {
Search(SearchResponse),
Image(ImageSearchResponse),
}
pub fn build_response_cache(ttl: Duration) -> Cache<String, CachedResponse> {
Cache::builder()
.time_to_live(ttl)
.max_capacity(10_000)
.build()
}
#[derive(Debug)]
pub struct EngineLimits {
semaphores: Mutex<HashMap<&'static str, Arc<Semaphore>>>,
cooldowns: Mutex<HashMap<&'static str, Instant>>,
max_concurrency: usize,
cooldown: Duration,
}
impl EngineLimits {
pub fn new(max_concurrency: usize, cooldown: Duration) -> Self {
Self {
semaphores: Mutex::new(HashMap::new()),
cooldowns: Mutex::new(HashMap::new()),
max_concurrency,
cooldown,
}
}
pub async fn acquire(&self, name: &'static str) -> Result<OwnedSemaphorePermit, EngineError> {
if self.in_cooldown(name) {
return Err(EngineError::Cooldown { engine: name });
}
let semaphore = self
.semaphores
.lock()
.expect("engine semaphore map poisoned")
.entry(name)
.or_insert_with(|| Arc::new(Semaphore::new(self.max_concurrency)))
.clone();
Ok(semaphore
.acquire_owned()
.await
.expect("engine semaphore closed"))
}
pub fn record_failure(&self, name: &'static str) {
self.cooldowns
.lock()
.expect("engine cooldown map poisoned")
.insert(name, Instant::now());
}
fn in_cooldown(&self, name: &'static str) -> bool {
let cooldowns = self.cooldowns.lock().expect("engine cooldown map poisoned");
matches!(
cooldowns.get(&name),
Some(started) if started.elapsed() < self.cooldown
)
}
}
impl Default for EngineLimits {
fn default() -> Self {
Self::new(4, Duration::from_secs(30))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn acquire_allows_concurrent_permits_up_to_the_limit() {
let limits = EngineLimits::new(2, Duration::from_secs(30));
let first = limits.acquire("engine").await.unwrap();
let second = limits.acquire("engine").await.unwrap();
let timed_out = tokio::time::timeout(Duration::from_millis(50), limits.acquire("engine"))
.await
.is_err();
assert!(timed_out);
drop(first);
let _third = tokio::time::timeout(Duration::from_millis(50), limits.acquire("engine"))
.await
.unwrap()
.unwrap();
drop(second);
}
#[tokio::test]
async fn record_failure_starts_cooldown() {
let limits = EngineLimits::new(2, Duration::from_secs(30));
limits.record_failure("engine");
assert!(matches!(
limits.acquire("engine").await,
Err(EngineError::Cooldown { engine: "engine" })
));
}
#[tokio::test]
async fn cooldown_expires_after_duration() {
let limits = EngineLimits::new(2, Duration::from_millis(20));
limits.record_failure("engine");
tokio::time::sleep(Duration::from_millis(50)).await;
let _ = limits.acquire("engine").await.unwrap();
}
#[tokio::test]
async fn cooldown_skips_do_not_extend_the_window() {
let limits = EngineLimits::new(2, Duration::from_millis(20));
limits.record_failure("engine");
for _ in 0..5 {
let _ = limits.acquire("engine").await;
tokio::time::sleep(Duration::from_millis(2)).await;
}
tokio::time::sleep(Duration::from_millis(20)).await;
let _ = limits.acquire("engine").await.unwrap();
}
}