aria2_core/selector/
server_stat_man.rs

1use serde::{Deserialize, Serialize};
2use std::collections::HashMap;
3use std::path::Path;
4use std::sync::{Arc, RwLock};
5
6use crate::selector::server_stat::{ServerStat, ServerStatSnapshot};
7
8/// Container for serialized server statistics file.
9///
10/// This is the top-level structure saved to disk when persisting
11/// server statistics. It includes metadata about when the file was
12/// saved and the version format.
13///
14/// # Example JSON structure
15///
16/// ```json
17/// {
18///   "version": "1.0",
19///   "saved_at": 1700000000,
20///   "servers": [
21///     {
22///       "host": "mirror.example.com",
23///       "download_speed": 5000000,
24///       ...
25///     }
26///   ]
27/// }
28/// ```
29#[derive(Debug, Clone, Serialize, Deserialize)]
30pub struct ServerStatFile {
31    /// File format version (currently "1.0")
32    pub version: String,
33    /// Unix timestamp when the file was saved
34    pub saved_at: u64,
35    /// List of server statistics
36    pub servers: Vec<ServerStatSnapshot>,
37}
38
39pub struct ServerStatMan {
40    stats: RwLock<HashMap<String, Arc<ServerStat>>>,
41}
42
43impl ServerStatMan {
44    pub fn new() -> Self {
45        Self {
46            stats: RwLock::new(HashMap::new()),
47        }
48    }
49
50    pub fn get_or_create(&self, host: &str) -> Arc<ServerStat> {
51        let mut map = self.stats.write().unwrap();
52        if let Some(stat) = map.get(host) {
53            Arc::clone(stat)
54        } else {
55            let stat = Arc::new(ServerStat::new(host));
56            map.insert(host.to_string(), Arc::clone(&stat));
57            stat
58        }
59    }
60
61    pub fn find_stat(&self, host: &str) -> Option<Arc<ServerStat>> {
62        let map = self.stats.read().unwrap();
63        map.get(host).cloned()
64    }
65
66    pub fn update(&self, host: &str, dl_speed: u64, is_multi: bool) {
67        let stat = self.get_or_create(host);
68        stat.update_speed(dl_speed, is_multi);
69    }
70
71    pub fn get_all_stats(&self) -> Vec<Arc<ServerStat>> {
72        let map = self.stats.read().unwrap();
73        map.values().cloned().collect()
74    }
75
76    pub fn remove(&self, host: &str) {
77        let mut map = self.stats.write().unwrap();
78        map.remove(host);
79    }
80
81    pub fn count(&self) -> usize {
82        let map = self.stats.read().unwrap();
83        map.len()
84    }
85
86    pub fn hosts(&self) -> Vec<String> {
87        let map = self.stats.read().unwrap();
88        map.keys().cloned().collect()
89    }
90
91    /// Mark a host as failed, updating error tracking fields.
92    ///
93    /// Clones the existing ServerStat, applies failure info via set_failure_info,
94    /// and replaces the entry in the map so all future Arc holders see the update.
95    pub fn mark_failure(&self, host: &str, error_code: u16) {
96        let mut map = self.stats.write().unwrap();
97        if let Some(stat_arc) = map.get(host) {
98            // Dereference Arc to get inner ServerStat, then clone the inner value
99            let inner: &ServerStat = stat_arc;
100            let mut updated = inner.clone();
101            updated.set_failure_info(error_code);
102            map.insert(host.to_string(), Arc::new(updated));
103        }
104    }
105
106    // ==================== Persistence Methods ====================
107
108    /// Save all server statistics to a JSON file (synchronous version).
109    ///
110    /// Serializes all server statistics to a JSON file at the specified path.
111    /// The file format includes version info, timestamp, and all server stats.
112    ///
113    /// # Arguments
114    ///
115    /// * `path` - File path to save to (e.g., "server-stat.json")
116    ///
117    /// # Returns
118    ///
119    /// * `Ok(usize)` - Number of servers saved
120    /// * `Err(String)` - Error message if save fails
121    ///
122    /// # Example
123    ///
124    /// ```no_run
125    /// use aria2_core::selector::server_stat_man::ServerStatMan;
126    /// use std::path::Path;
127    ///
128    /// let man = ServerStatMan::new();
129    /// man.update("mirror.example.com", 5000, false);
130    ///
131    /// let saved = man.save_to_file(Path::new("server-stat.json")).unwrap();
132    /// println!("Saved {} servers", saved);
133    /// ```
134    pub fn save_to_file(&self, path: &Path) -> Result<usize, String> {
135        let map = self.stats.read().unwrap();
136
137        let servers: Vec<ServerStatSnapshot> =
138            map.values().map(|stat| stat.to_snapshot()).collect();
139
140        let file_content = ServerStatFile {
141            version: "1.0".to_string(),
142            saved_at: std::time::SystemTime::now()
143                .duration_since(std::time::UNIX_EPOCH)
144                .unwrap_or_default()
145                .as_secs(),
146            servers,
147        };
148
149        let json = serde_json::to_string_pretty(&file_content)
150            .map_err(|e| format!("Failed to serialize server stats: {}", e))?;
151
152        std::fs::write(path, json)
153            .map_err(|e| format!("Failed to write server stat file: {}", e))?;
154
155        Ok(file_content.servers.len())
156    }
157
158    /// Load server statistics from a JSON file (synchronous version).
159    ///
160    /// Reads a JSON file and restores all server statistics into memory.
161    /// Existing statistics are preserved; loaded stats are merged in.
162    ///
163    /// # Arguments
164    ///
165    /// * `path` - File path to load from
166    ///
167    /// # Returns
168    ///
169    /// * `Ok(usize)` - Number of servers loaded
170    /// * `Err(String)` - Error message if load fails
171    ///
172    /// # Behavior
173    ///
174    /// - Returns `Ok(0)` if the file doesn't exist (not an error)
175    /// - Returns error if file exists but is invalid JSON
176    /// - Merges with existing stats (doesn't clear current stats)
177    ///
178    /// # Example
179    ///
180    /// ```no_run
181    /// use aria2_core::selector::server_stat_man::ServerStatMan;
182    /// use std::path::Path;
183    ///
184    /// let man = ServerStatMan::new();
185    /// let loaded = man.load_from_file(Path::new("server-stat.json")).unwrap();
186    /// println!("Loaded {} servers", loaded);
187    /// ```
188    pub fn load_from_file(&self, path: &Path) -> Result<usize, String> {
189        if !path.exists() {
190            return Ok(0); // Not an error if file doesn't exist
191        }
192
193        let content = std::fs::read_to_string(path)
194            .map_err(|e| format!("Failed to read server stat file: {}", e))?;
195
196        let file_content: ServerStatFile = serde_json::from_str(&content)
197            .map_err(|e| format!("Invalid server stat file format: {}", e))?;
198
199        let count = file_content.servers.len();
200        let mut map = self.stats.write().unwrap();
201
202        for snapshot in file_content.servers {
203            let stat = Arc::new(ServerStat::from_snapshot(&snapshot));
204            map.insert(snapshot.host, stat);
205        }
206
207        Ok(count)
208    }
209
210    /// Save all server statistics to a JSON file (async version).
211    ///
212    /// Async variant of `save_to_file()` for use in async contexts.
213    /// This is the preferred method when called from async code.
214    ///
215    /// # Arguments
216    ///
217    /// * `path` - File path to save to
218    ///
219    /// # Returns
220    ///
221    /// * `Ok(usize)` - Number of servers saved
222    /// * `Err(String)` - Error message if save fails
223    pub async fn save_to_file_async(&self, path: &Path) -> Result<usize, String> {
224        let path = path.to_path_buf();
225        let snapshots: Vec<ServerStatSnapshot> = {
226            let map = self.stats.read().unwrap();
227            map.values().map(|stat| stat.to_snapshot()).collect()
228        };
229
230        let file_content = ServerStatFile {
231            version: "1.0".to_string(),
232            saved_at: std::time::SystemTime::now()
233                .duration_since(std::time::UNIX_EPOCH)
234                .unwrap_or_default()
235                .as_secs(),
236            servers: snapshots,
237        };
238
239        let json = serde_json::to_string_pretty(&file_content)
240            .map_err(|e| format!("Failed to serialize server stats: {}", e))?;
241
242        tokio::fs::write(&path, json)
243            .await
244            .map_err(|e| format!("Failed to write server stat file: {}", e))?;
245
246        Ok(file_content.servers.len())
247    }
248
249    /// Load server statistics from a JSON file (async version).
250    ///
251    /// Async variant of `load_from_file()` for use in async contexts.
252    /// This is the preferred method when called from async code.
253    ///
254    /// # Arguments
255    ///
256    /// * `path` - File path to load from
257    ///
258    /// # Returns
259    ///
260    /// * `Ok(usize)` - Number of servers loaded
261    /// * `Err(String)` - Error message if load fails
262    pub async fn load_from_file_async(&self, path: &Path) -> Result<usize, String> {
263        if !path.exists() {
264            return Ok(0);
265        }
266
267        let content = tokio::fs::read_to_string(path)
268            .await
269            .map_err(|e| format!("Failed to read server stat file: {}", e))?;
270
271        let file_content: ServerStatFile = serde_json::from_str(&content)
272            .map_err(|e| format!("Invalid server stat file format: {}", e))?;
273
274        let count = file_content.servers.len();
275        let mut map = self.stats.write().unwrap();
276        for snapshot in file_content.servers {
277            let stat = Arc::new(ServerStat::from_snapshot(&snapshot));
278            map.insert(snapshot.host, stat);
279        }
280
281        Ok(count)
282    }
283}
284
285impl Default for ServerStatMan {
286    fn default() -> Self {
287        Self::new()
288    }
289}
290
291#[cfg(test)]
292mod tests {
293    use super::*;
294
295    #[test]
296    fn test_creation_and_count() {
297        let man = ServerStatMan::new();
298        assert_eq!(man.count(), 0);
299    }
300
301    #[test]
302    fn test_get_or_create_new_host() {
303        let man = ServerStatMan::new();
304        let stat = man.get_or_create("example.com");
305        assert_eq!(stat.host, "example.com");
306        assert_eq!(man.count(), 1);
307    }
308
309    #[test]
310    fn test_get_or_create_returns_same_instance() {
311        let man = ServerStatMan::new();
312        let s1 = man.get_or_create("example.com");
313        let s2 = man.get_or_create("example.com");
314        assert!(Arc::ptr_eq(&s1, &s2));
315        assert_eq!(man.count(), 1);
316    }
317
318    #[test]
319    fn test_find_existing() {
320        let man = ServerStatMan::new();
321        man.get_or_create("example.com");
322        assert!(man.find_stat("example.com").is_some());
323        assert!(man.find_stat("nonexistent").is_none());
324    }
325
326    #[test]
327    fn test_update_creates_if_needed() {
328        let man = ServerStatMan::new();
329        man.update("fast.server", 5000, false);
330        assert_eq!(man.count(), 1);
331
332        let stat = man.find_stat("fast.server").unwrap();
333        assert_eq!(stat.get_download_speed(), 5000);
334    }
335
336    #[test]
337    fn test_remove() {
338        let man = ServerStatMan::new();
339        man.get_or_create("a.com");
340        man.get_or_create("b.com");
341        assert_eq!(man.count(), 2);
342
343        man.remove("a.com");
344        assert_eq!(man.count(), 1);
345        assert!(man.find_stat("a.com").is_none());
346    }
347
348    #[test]
349    fn test_multiple_hosts_independent() {
350        let man = ServerStatMan::new();
351        man.update("fast.com", 10000, true);
352        man.update("slow.com", 100, false);
353
354        let fast = man.find_stat("fast.com").unwrap();
355        let slow = man.find_stat("slow.com").unwrap();
356
357        assert_ne!(fast.get_avg_speed(), slow.get_avg_speed());
358        assert!(fast.get_avg_speed() > slow.get_avg_speed());
359    }
360
361    #[test]
362    fn test_hosts_list() {
363        let man = ServerStatMan::new();
364        man.get_or_create("alpha.com");
365        man.get_or_create("beta.com");
366        let hosts = man.hosts();
367        assert_eq!(hosts.len(), 2);
368        assert!(hosts.contains(&"alpha.com".to_string()));
369        assert!(hosts.contains(&"beta.com".to_string()));
370    }
371
372    // ======================================================================
373    // Tests for mark_failure
374    // ======================================================================
375
376    #[test]
377    fn test_mark_failure_updates_stats() {
378        let man = ServerStatMan::new();
379        man.get_or_create("failing.host");
380
381        // Mark as failed with error code 500
382        man.mark_failure("failing.host", 500);
383
384        let stat = man.find_stat("failing.host").unwrap();
385        assert_eq!(stat.get_consecutive_failures(), 1);
386        assert!(stat.get_last_error_time() > 0);
387        assert_eq!(stat.get_last_error_code(), 500);
388    }
389
390    #[test]
391    fn test_mark_failure_multiple_times() {
392        let man = ServerStatMan::new();
393        man.get_or_create("repeated.failures");
394
395        for i in 0..5u16 {
396            man.mark_failure("repeated.failures", i);
397        }
398
399        let stat = man.find_stat("repeated.failures").unwrap();
400        assert_eq!(stat.get_consecutive_failures(), 5);
401        assert!(
402            !stat.is_available(),
403            "Should be unavailable after 5 failures"
404        );
405    }
406
407    #[test]
408    fn test_mark_failure_nonexistent_host() {
409        let man = ServerStatMan::new();
410        // Should not panic on nonexistent host
411        man.mark_failure("nonexistent.host", 404);
412        assert_eq!(man.count(), 0); // No stat created
413    }
414
415    // ======================================================================
416    // Tests for Persistence (save/load)
417    // ======================================================================
418
419    #[test]
420    fn test_save_to_file_basic() {
421        let man = ServerStatMan::new();
422        man.update("fast.mirror.com", 10000, false);
423        man.update("slow.mirror.com", 100, true);
424
425        let temp_file = std::env::temp_dir().join("test_server_stat_save.json");
426        let saved = man.save_to_file(&temp_file).expect("Save should succeed");
427
428        assert_eq!(saved, 2, "Should save 2 servers");
429
430        // Verify file content is valid JSON
431        let content = std::fs::read_to_string(&temp_file).expect("Should read file");
432        let parsed: serde_json::Value =
433            serde_json::from_str(&content).expect("Should be valid JSON");
434        assert!(parsed.get("version").is_some());
435        assert!(parsed.get("saved_at").is_some());
436        assert!(parsed.get("servers").is_some());
437
438        let _ = std::fs::remove_file(&temp_file);
439    }
440
441    #[test]
442    fn test_load_from_file_basic() {
443        let man = ServerStatMan::new();
444        man.update("load.test.com", 5000, false);
445        man.mark_failure("load.test.com", 500);
446
447        let temp_file = std::env::temp_dir().join("test_server_stat_load.json");
448        let _ = man.save_to_file(&temp_file);
449
450        // Load into a new manager
451        let man2 = ServerStatMan::new();
452        let loaded = man2
453            .load_from_file(&temp_file)
454            .expect("Load should succeed");
455
456        assert_eq!(loaded, 1, "Should load 1 server");
457
458        let stat = man2
459            .find_stat("load.test.com")
460            .expect("Should find loaded server");
461        assert!(stat.get_avg_speed() > 0);
462        assert_eq!(stat.get_consecutive_failures(), 1);
463
464        let _ = std::fs::remove_file(&temp_file);
465    }
466
467    #[test]
468    fn test_save_load_roundtrip() {
469        let man = ServerStatMan::new();
470        man.update("mirror1.com", 10000, false);
471        man.update("mirror2.com", 5000, true);
472        man.update("mirror3.com", 2000, false);
473        man.mark_failure("mirror3.com", 503);
474        man.mark_failure("mirror3.com", 503);
475
476        let temp_file = std::env::temp_dir().join("test_server_stat_roundtrip.json");
477        let saved = man.save_to_file(&temp_file).unwrap();
478        assert_eq!(saved, 3);
479
480        let man2 = ServerStatMan::new();
481        let loaded = man2.load_from_file(&temp_file).unwrap();
482        assert_eq!(loaded, 3);
483
484        // Verify all servers restored correctly
485        let s1 = man2.find_stat("mirror1.com").unwrap();
486        let s2 = man2.find_stat("mirror2.com").unwrap();
487        let s3 = man2.find_stat("mirror3.com").unwrap();
488
489        assert!(s1.get_single_avg_speed() > 0);
490        assert!(s2.get_multi_avg_speed() > 0);
491        assert_eq!(s3.get_consecutive_failures(), 2);
492        assert_eq!(s3.get_last_error_code(), 503);
493
494        let _ = std::fs::remove_file(&temp_file);
495    }
496
497    #[test]
498    fn test_load_nonexistent_file() {
499        let man = ServerStatMan::new();
500        let nonexistent = std::path::Path::new("/tmp/nonexistent_server_stat_12345.json");
501
502        let loaded = man
503            .load_from_file(nonexistent)
504            .expect("Should return Ok(0) for nonexistent file");
505        assert_eq!(loaded, 0);
506        assert_eq!(man.count(), 0);
507    }
508
509    #[test]
510    fn test_save_empty_manager() {
511        let man = ServerStatMan::new();
512        let temp_file = std::env::temp_dir().join("test_server_stat_empty.json");
513
514        let saved = man
515            .save_to_file(&temp_file)
516            .expect("Should save empty manager");
517        assert_eq!(saved, 0);
518
519        // Verify file is still valid JSON
520        let content = std::fs::read_to_string(&temp_file).expect("Should read file");
521        let parsed: serde_json::Value =
522            serde_json::from_str(&content).expect("Should be valid JSON");
523        let servers = parsed.get("servers").unwrap().as_array().unwrap();
524        assert!(servers.is_empty());
525
526        let _ = std::fs::remove_file(&temp_file);
527    }
528
529    #[test]
530    fn test_load_merges_with_existing() {
531        let man = ServerStatMan::new();
532        man.update("existing.com", 1000, false);
533
534        // Create a file with a different server
535        let temp_file = std::env::temp_dir().join("test_server_stat_merge.json");
536        let man2 = ServerStatMan::new();
537        man2.update("fromfile.com", 5000, false);
538        let _ = man2.save_to_file(&temp_file);
539
540        // Load into man (which already has existing.com)
541        let loaded = man.load_from_file(&temp_file).unwrap();
542        assert_eq!(loaded, 1);
543
544        // Should have both servers now
545        assert_eq!(man.count(), 2);
546        assert!(man.find_stat("existing.com").is_some());
547        assert!(man.find_stat("fromfile.com").is_some());
548
549        let _ = std::fs::remove_file(&temp_file);
550    }
551
552    #[tokio::test]
553    async fn test_async_save_and_load() {
554        let man = ServerStatMan::new();
555        man.update("async.test.com", 8000, false);
556        man.update("async2.test.com", 4000, true);
557
558        let temp_file = std::env::temp_dir().join("test_server_stat_async.json");
559
560        // Async save
561        let saved = man
562            .save_to_file_async(&temp_file)
563            .await
564            .expect("Async save should succeed");
565        assert_eq!(saved, 2);
566
567        // Async load
568        let man2 = ServerStatMan::new();
569        let loaded = man2
570            .load_from_file_async(&temp_file)
571            .await
572            .expect("Async load should succeed");
573        assert_eq!(loaded, 2);
574
575        let stat = man2.find_stat("async.test.com").unwrap();
576        assert!(stat.get_avg_speed() > 0);
577
578        let _ = std::fs::remove_file(&temp_file);
579    }
580
581    #[test]
582    fn test_file_format_structure() {
583        let man = ServerStatMan::new();
584        man.update("format.test", 12345, false);
585
586        let temp_file = std::env::temp_dir().join("test_server_stat_format.json");
587        let _ = man.save_to_file(&temp_file);
588
589        let content = std::fs::read_to_string(&temp_file).unwrap();
590        let parsed: ServerStatFile = serde_json::from_str(&content).unwrap();
591
592        assert_eq!(parsed.version, "1.0");
593        assert!(parsed.saved_at > 0);
594        assert_eq!(parsed.servers.len(), 1);
595        assert_eq!(parsed.servers[0].host, "format.test");
596        assert_eq!(parsed.servers[0].download_speed, 12345);
597
598        let _ = std::fs::remove_file(&temp_file);
599    }
600}
aria2_core/selector/server_stat_man.rs

aria2_core/selector/
server_stat_man.rs
