Skip to main content

turbo_cdn/
dns_cache.rs

1// Licensed under the MIT License
2// Copyright (c) 2025 Hal <hal.long@outlook.com>
3
4//! DNS Cache System
5//!
6//! This module implements a high-performance DNS cache to reduce DNS query latency
7//! and improve overall download performance.
8
9use 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/// DNS cache entry
17#[derive(Debug, Clone)]
18pub struct DnsCacheEntry {
19    /// Resolved IP addresses
20    pub addresses: Vec<IpAddr>,
21    /// Time when entry was created
22    pub created_at: Instant,
23    /// Time-to-live for this entry
24    pub ttl: Duration,
25    /// Number of times this entry has been used
26    pub hit_count: u64,
27    /// Last time this entry was accessed
28    pub last_accessed: Instant,
29}
30
31impl DnsCacheEntry {
32    /// Check if this cache entry is still valid
33    pub fn is_valid(&self) -> bool {
34        self.created_at.elapsed() < self.ttl
35    }
36
37    /// Check if this entry is expired
38    pub fn is_expired(&self) -> bool {
39        !self.is_valid()
40    }
41
42    /// Update access statistics
43    pub fn mark_accessed(&mut self) {
44        self.hit_count += 1;
45        self.last_accessed = Instant::now();
46    }
47}
48
49/// DNS cache statistics
50#[derive(Debug, Clone)]
51pub struct DnsCacheStats {
52    /// Total number of cache entries
53    pub total_entries: usize,
54    /// Number of cache hits
55    pub cache_hits: u64,
56    /// Number of cache misses
57    pub cache_misses: u64,
58    /// Cache hit ratio
59    pub hit_ratio: f64,
60    /// Number of expired entries cleaned up
61    pub expired_cleaned: u64,
62    /// Total DNS resolution time saved (estimated)
63    pub time_saved_ms: u64,
64}
65
66/// High-performance DNS cache
67#[derive(Debug)]
68pub struct DnsCache {
69    /// Cache storage using DashMap for concurrent access
70    cache: DashMap<String, DnsCacheEntry>,
71    /// Cache statistics
72    stats: Arc<RwLock<DnsCacheStats>>,
73    /// Default TTL for cache entries
74    default_ttl: Duration,
75    /// Maximum number of cache entries
76    max_entries: usize,
77    /// Cleanup interval
78    cleanup_interval: Duration,
79    /// Last cleanup time
80    last_cleanup: Arc<RwLock<Instant>>,
81}
82
83impl DnsCache {
84    /// Create a new DNS cache
85    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), // 5 minutes
99            last_cleanup: Arc::new(RwLock::new(Instant::now())),
100        }
101    }
102
103    /// Create a DNS cache with default settings
104    pub fn with_defaults() -> Self {
105        Self::new(
106            Duration::from_secs(300), // 5 minutes TTL
107            1000,                     // Max 1000 entries
108        )
109    }
110
111    /// Resolve hostname with caching
112    pub async fn resolve(&self, hostname: &str) -> Option<Vec<IpAddr>> {
113        // Check cache first
114        if let Some(mut entry) = self.cache.get_mut(hostname) {
115            if entry.is_valid() {
116                entry.mark_accessed();
117
118                // Update statistics
119                let mut stats = self.stats.write().await;
120                stats.cache_hits += 1;
121                stats.time_saved_ms += 50; // Estimate 50ms saved per cache hit
122                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                // Entry expired, remove it
133                drop(entry);
134                self.cache.remove(hostname);
135                debug!("Removed expired DNS cache entry for {}", hostname);
136            }
137        }
138
139        // Cache miss - perform actual DNS resolution
140        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                    // Update statistics
146                    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                // Update statistics
166                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    /// Insert entry into cache
177    pub async fn insert(&self, hostname: &str, addresses: Vec<IpAddr>, ttl: Option<Duration>) {
178        // Check if cleanup is needed
179        self.maybe_cleanup().await;
180
181        // Enforce max entries limit
182        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        // Update statistics
197        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    /// Perform actual DNS resolution
208    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    /// Update cache hit ratio
227    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    /// Maybe perform cache cleanup
237    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    /// Clean up expired cache entries
246    pub async fn cleanup_expired(&self) {
247        let mut expired_count = 0;
248        let mut to_remove = Vec::new();
249
250        // Collect expired entries
251        for entry in self.cache.iter() {
252            if entry.value().is_expired() {
253                to_remove.push(entry.key().clone());
254            }
255        }
256
257        // Remove expired entries
258        for key in to_remove {
259            self.cache.remove(&key);
260            expired_count += 1;
261        }
262
263        if expired_count > 0 {
264            // Update statistics
265            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    /// Evict oldest entries when cache is full
275    async fn evict_oldest_entries(&self) {
276        let target_size = self.max_entries * 3 / 4; // Remove 25% of entries
277        let mut entries_to_remove = Vec::new();
278
279        // Collect entries with their last access time
280        let mut entries: Vec<_> = self
281            .cache
282            .iter()
283            .map(|entry| (entry.key().clone(), entry.value().last_accessed))
284            .collect();
285
286        // Sort by last accessed time (oldest first)
287        entries.sort_by_key(|(_, last_accessed)| *last_accessed);
288
289        // Mark oldest entries for removal
290        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        // Remove entries
296        for key in entries_to_remove {
297            self.cache.remove(&key);
298        }
299
300        // Update statistics
301        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    /// Get cache statistics
309    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    /// Clear all cache entries
316    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    /// Get cache entry for hostname (for debugging)
327    pub fn get_entry(&self, hostname: &str) -> Option<DnsCacheEntry> {
328        self.cache.get(hostname).map(|entry| entry.clone())
329    }
330
331    /// Check if hostname is cached
332    pub fn contains(&self, hostname: &str) -> bool {
333        self.cache.contains_key(hostname)
334    }
335
336    /// Get number of cached entries
337    pub fn len(&self) -> usize {
338        self.cache.len()
339    }
340
341    /// Check if cache is empty
342    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}