diff --git a/net-guardia/src/core/infrastructure/geoip.rs b/net-guardia/src/core/infrastructure/geoip.rs index a1f39b1..0bca797 100644 --- a/net-guardia/src/core/infrastructure/geoip.rs +++ b/net-guardia/src/core/infrastructure/geoip.rs @@ -3,7 +3,7 @@ use std::path::{Path, PathBuf}; use std::sync::Arc; use maxminddb::{geoip2, MaxMindDbError, Reader}; -use tokio::sync::RwLock; +use tokio::sync::{RwLock, Semaphore}; use lru::LruCache; use std::num::NonZeroUsize; use tokio::task; @@ -11,9 +11,12 @@ use tokio::task; use crate::model::geo_stats::GeoLocation; use crate::utils::ip_address; +const MAX_CONCURRENT_DB_LOOKUPS: usize = 8; + pub struct GeoIpService { reader: Arc>>, cache: Arc>>>, + lookup_sem: Arc, } impl GeoIpService { @@ -33,6 +36,7 @@ impl GeoIpService { Ok(Self { reader: Arc::new(reader), cache: Arc::new(RwLock::new(LruCache::new(cache_capacity))), + lookup_sem: Arc::new(Semaphore::new(MAX_CONCURRENT_DB_LOOKUPS)), }) } @@ -55,8 +59,15 @@ impl GeoIpService { } } + let permit = self.lookup_sem.clone().acquire_owned().await + .map_err(|_| MaxMindDbError::InvalidDatabase { + message: "GeoIP semaphore closed".to_string(), + offset: None, + })?; + let reader = self.reader.clone(); let result = task::spawn_blocking(move || { + let _permit = permit; Self::lookup_from_db_blocking(&reader, ip) }) .await