rust_nrm/utils/
speed_test.rs

1use super::Logger;
2use colored::Colorize;
3use reqwest;
4use std::sync::atomic::{AtomicUsize, Ordering};
5use std::sync::Arc;
6use std::time::Instant;
7use tokio::sync::Mutex;
8use tokio::task;
9
10#[derive(Debug, Clone)]
11pub struct SpeedTestResult {
12    pub name: String,
13    pub url: String,
14    pub response_time: f64,
15    pub is_success: bool,
16}
17
18#[derive(Clone)]
19pub struct SpeedTester {
20    client: reqwest::Client,
21}
22
23impl SpeedTester {
24    pub fn new() -> Self {
25        SpeedTester {
26            client: reqwest::Client::builder()
27                .timeout(std::time::Duration::from_secs(5)) // 5 seconds timeout
28                .build()
29                .unwrap_or_else(|_| reqwest::Client::new()),
30        }
31    }
32
33    pub async fn test_registry(&self, name: &str, url: &str) -> SpeedTestResult {
34        let start = Instant::now();
35        let result = self.client.get(url).send().await;
36        let elapsed = start.elapsed().as_secs_f64();
37
38        SpeedTestResult {
39            name: name.to_string(),
40            url: url.to_string(),
41            response_time: elapsed,
42            is_success: result.is_ok(),
43        }
44    }
45
46    fn format_time(seconds: f64) -> String {
47        format!("{:.0}ms", seconds * 1000.0)
48    }
49
50    pub async fn test_all(&self, registries: &[(String, String)]) -> Vec<SpeedTestResult> {
51        Logger::info("Testing registry speeds...");
52        println!(); // Add a blank line for better readability
53
54        // Create a shared vector for results
55        let results = Arc::new(Mutex::new(Vec::new()));
56
57        // Create an atomic counter for ranking
58        let speed_rank = Arc::new(AtomicUsize::new(1));
59
60        // Create a vector of tasks for parallel execution
61        let mut handles: Vec<task::JoinHandle<()>> = Vec::new();
62        for (name, url) in registries.iter() {
63            let name = name.clone();
64            let url = url.clone();
65            let client = self.clone();
66            let results = Arc::clone(&results);
67            let speed_rank = Arc::clone(&speed_rank);
68
69            let handle = task::spawn(async move {
70                let test_result = client.test_registry(&name, &url).await;
71
72                let status = if test_result.is_success {
73                    "✓".green()
74                } else {
75                    "✗".red()
76                };
77
78                let time_str = Self::format_time(test_result.response_time);
79                let time_display = time_str.normal();
80                let current_rank = speed_rank.fetch_add(1, Ordering::Relaxed);
81                let rank = format!("#{}", current_rank).bold();
82
83                println!(
84                    "{} {} {} {} {}",
85                    rank,
86                    status,
87                    test_result.name.bold(),
88                    "->".dimmed(),
89                    time_display
90                );
91
92                // Store the result
93                let mut results = results.lock().await;
94                results.push(test_result);
95            });
96            handles.push(handle);
97        }
98
99        // Wait for all tasks to complete
100        for handle in handles {
101            let _ = handle.await;
102        }
103
104        // Get all results
105        let final_results = results.lock().await.clone();
106
107        // Find the fastest successful registry
108        let fastest = final_results
109            .iter()
110            .filter(|r| r.is_success)
111            .min_by(|a, b| a.response_time.partial_cmp(&b.response_time).unwrap());
112
113        // Show the fastest registry
114        if let Some(fastest) = fastest {
115            println!(); // Add a blank line
116            Logger::success(&format!(
117                "Fastest registry is {} ({})",
118                fastest.name.bold(),
119                Self::format_time(fastest.response_time)
120            ));
121        }
122
123        final_results
124    }
125}
rust_nrm/utils/speed_test.rs

rust_nrm/utils/
speed_test.rs
