1use dashmap::DashMap;
10use std::net::IpAddr;
11use std::sync::Arc;
12use std::time::{Duration, Instant};
13use tokio::sync::RwLock;
14use tracing::{debug, info, warn};
15
16#[derive(Debug, Clone)]
18pub struct DnsCacheEntry {
19 pub addresses: Vec<IpAddr>,
21 pub created_at: Instant,
23 pub ttl: Duration,
25 pub hit_count: u64,
27 pub last_accessed: Instant,
29}
30
31impl DnsCacheEntry {
32 pub fn is_valid(&self) -> bool {
34 self.created_at.elapsed() < self.ttl
35 }
36
37 pub fn is_expired(&self) -> bool {
39 !self.is_valid()
40 }
41
42 pub fn mark_accessed(&mut self) {
44 self.hit_count += 1;
45 self.last_accessed = Instant::now();
46 }
47}
48
49#[derive(Debug, Clone)]
51pub struct DnsCacheStats {
52 pub total_entries: usize,
54 pub cache_hits: u64,
56 pub cache_misses: u64,
58 pub hit_ratio: f64,
60 pub expired_cleaned: u64,
62 pub time_saved_ms: u64,
64}
65
66#[derive(Debug)]
68pub struct DnsCache {
69 cache: DashMap<String, DnsCacheEntry>,
71 stats: Arc<RwLock<DnsCacheStats>>,
73 default_ttl: Duration,
75 max_entries: usize,
77 cleanup_interval: Duration,
79 last_cleanup: Arc<RwLock<Instant>>,
81}
82
83impl DnsCache {
84 pub fn new(default_ttl: Duration, max_entries: usize) -> Self {
86 Self {
87 cache: DashMap::new(),
88 stats: Arc::new(RwLock::new(DnsCacheStats {
89 total_entries: 0,
90 cache_hits: 0,
91 cache_misses: 0,
92 hit_ratio: 0.0,
93 expired_cleaned: 0,
94 time_saved_ms: 0,
95 })),
96 default_ttl,
97 max_entries,
98 cleanup_interval: Duration::from_secs(300), last_cleanup: Arc::new(RwLock::new(Instant::now())),
100 }
101 }
102
103 pub fn with_defaults() -> Self {
105 Self::new(
106 Duration::from_secs(300), 1000, )
109 }
110
111 pub async fn resolve(&self, hostname: &str) -> Option<Vec<IpAddr>> {
113 if let Some(mut entry) = self.cache.get_mut(hostname) {
115 if entry.is_valid() {
116 entry.mark_accessed();
117
118 let mut stats = self.stats.write().await;
120 stats.cache_hits += 1;
121 stats.time_saved_ms += 50; self.update_hit_ratio(&mut stats);
123 drop(stats);
124
125 debug!(
126 "DNS cache hit for {}: {} addresses",
127 hostname,
128 entry.addresses.len()
129 );
130 return Some(entry.addresses.clone());
131 } else {
132 drop(entry);
134 self.cache.remove(hostname);
135 debug!("Removed expired DNS cache entry for {}", hostname);
136 }
137 }
138
139 match self.perform_dns_resolution(hostname).await {
141 Ok(addresses) => {
142 if !addresses.is_empty() {
143 self.insert(hostname, addresses.clone(), None).await;
144
145 let mut stats = self.stats.write().await;
147 stats.cache_misses += 1;
148 self.update_hit_ratio(&mut stats);
149 drop(stats);
150
151 debug!(
152 "DNS resolved and cached for {}: {} addresses",
153 hostname,
154 addresses.len()
155 );
156 Some(addresses)
157 } else {
158 warn!("DNS resolution returned no addresses for {}", hostname);
159 None
160 }
161 }
162 Err(e) => {
163 warn!("DNS resolution failed for {}: {}", hostname, e);
164
165 let mut stats = self.stats.write().await;
167 stats.cache_misses += 1;
168 self.update_hit_ratio(&mut stats);
169 drop(stats);
170
171 None
172 }
173 }
174 }
175
176 pub async fn insert(&self, hostname: &str, addresses: Vec<IpAddr>, ttl: Option<Duration>) {
178 self.maybe_cleanup().await;
180
181 if self.cache.len() >= self.max_entries {
183 self.evict_oldest_entries().await;
184 }
185
186 let entry = DnsCacheEntry {
187 addresses,
188 created_at: Instant::now(),
189 ttl: ttl.unwrap_or(self.default_ttl),
190 hit_count: 0,
191 last_accessed: Instant::now(),
192 };
193
194 self.cache.insert(hostname.to_string(), entry);
195
196 let mut stats = self.stats.write().await;
198 stats.total_entries = self.cache.len();
199 drop(stats);
200
201 debug!(
202 "Cached DNS entry for {} (TTL: {:?})",
203 hostname, self.default_ttl
204 );
205 }
206
207 async fn perform_dns_resolution(
209 &self,
210 hostname: &str,
211 ) -> Result<Vec<IpAddr>, Box<dyn std::error::Error + Send + Sync>> {
212 use tokio::net::lookup_host;
213
214 let start = Instant::now();
215 let addrs: Vec<IpAddr> = lookup_host((hostname, 80))
216 .await?
217 .map(|addr| addr.ip())
218 .collect();
219
220 let duration = start.elapsed();
221 debug!("DNS resolution for {} took {:?}", hostname, duration);
222
223 Ok(addrs)
224 }
225
226 fn update_hit_ratio(&self, stats: &mut DnsCacheStats) {
228 let total = stats.cache_hits + stats.cache_misses;
229 stats.hit_ratio = if total > 0 {
230 stats.cache_hits as f64 / total as f64
231 } else {
232 0.0
233 };
234 }
235
236 async fn maybe_cleanup(&self) {
238 let last_cleanup = *self.last_cleanup.read().await;
239 if last_cleanup.elapsed() >= self.cleanup_interval {
240 self.cleanup_expired().await;
241 *self.last_cleanup.write().await = Instant::now();
242 }
243 }
244
245 pub async fn cleanup_expired(&self) {
247 let mut expired_count = 0;
248 let mut to_remove = Vec::new();
249
250 for entry in self.cache.iter() {
252 if entry.value().is_expired() {
253 to_remove.push(entry.key().clone());
254 }
255 }
256
257 for key in to_remove {
259 self.cache.remove(&key);
260 expired_count += 1;
261 }
262
263 if expired_count > 0 {
264 let mut stats = self.stats.write().await;
266 stats.expired_cleaned += expired_count;
267 stats.total_entries = self.cache.len();
268 drop(stats);
269
270 info!("Cleaned up {} expired DNS cache entries", expired_count);
271 }
272 }
273
274 async fn evict_oldest_entries(&self) {
276 let target_size = self.max_entries * 3 / 4; let mut entries_to_remove = Vec::new();
278
279 let mut entries: Vec<_> = self
281 .cache
282 .iter()
283 .map(|entry| (entry.key().clone(), entry.value().last_accessed))
284 .collect();
285
286 entries.sort_by_key(|(_, last_accessed)| *last_accessed);
288
289 let remove_count = self.cache.len().saturating_sub(target_size);
291 for (key, _) in entries.into_iter().take(remove_count) {
292 entries_to_remove.push(key);
293 }
294
295 for key in entries_to_remove {
297 self.cache.remove(&key);
298 }
299
300 let mut stats = self.stats.write().await;
302 stats.total_entries = self.cache.len();
303 drop(stats);
304
305 info!("Evicted {} oldest DNS cache entries", remove_count);
306 }
307
308 pub async fn get_stats(&self) -> DnsCacheStats {
310 let mut stats = self.stats.read().await.clone();
311 stats.total_entries = self.cache.len();
312 stats
313 }
314
315 pub async fn clear(&self) {
317 self.cache.clear();
318
319 let mut stats = self.stats.write().await;
320 stats.total_entries = 0;
321 drop(stats);
322
323 info!("DNS cache cleared");
324 }
325
326 pub fn get_entry(&self, hostname: &str) -> Option<DnsCacheEntry> {
328 self.cache.get(hostname).map(|entry| entry.clone())
329 }
330
331 pub fn contains(&self, hostname: &str) -> bool {
333 self.cache.contains_key(hostname)
334 }
335
336 pub fn len(&self) -> usize {
338 self.cache.len()
339 }
340
341 pub fn is_empty(&self) -> bool {
343 self.cache.is_empty()
344 }
345}
346
347impl std::fmt::Display for DnsCacheStats {
348 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
349 write!(f,
350 "DNS Cache: {} entries | Hit ratio: {:.1}% ({}/{}) | Time saved: {}ms | Expired cleaned: {}",
351 self.total_entries,
352 self.hit_ratio * 100.0,
353 self.cache_hits,
354 self.cache_hits + self.cache_misses,
355 self.time_saved_ms,
356 self.expired_cleaned
357 )
358 }
359}