Untitled

use std::net::{IpAddr, Ipv4Addr, SocketAddr};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use rand::{rngs::StdRng, Rng, SeedableRng};
use trust_dns_resolver::config::{NameServerConfig, Protocol, ResolverConfig, ResolverOpts};
use trust_dns_resolver::TokioAsyncResolver;

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    // 配置DNS服务器
    let dns_server_ip = "172.17.17.1".parse::<Ipv4Addr>()?;
    
    // 创建解析器配置
    let mut config = ResolverConfig::new();
    let nameserver_config = NameServerConfig {
        socket_addr: SocketAddr::new(IpAddr::V4(dns_server_ip), 53),
        protocol: Protocol::Udp,
        tls_dns_name: None,
        trust_nx_responses: false,
    };
    config.add_name_server(nameserver_config);

    // 初始化异步解析器
    let resolver = TokioAsyncResolver::tokio(config, ResolverOpts::default())?;

    // 定义测试域名列表
    let domains = vec![
        "www.youtube.com",
        "www.google.com",
        "www.hub.docker.com",
        "www.pronhub.com",
    ];

    // 创建性能计数器
    let success = Arc::new(AtomicU64::new(0));
    let failure = Arc::new(AtomicU64::new(0));

    // 启动统计输出任务
    let stats_success = success.clone();
    let stats_failure = failure.clone();
    tokio::spawn(async move {
        loop {
            tokio::time::sleep(std::time::Duration::from_secs(1)).await;
            let s = stats_success.load(Ordering::Relaxed);
            let f = stats_failure.load(Ordering::Relaxed);
            let total = s + f;
            println!(
                "Requests: {} | Success: {} ({:.2}%) | Failures: {}",
                total,
                s,
                if total > 0 { s as f64 / total as f64 * 100.0 } else { 0.0 },
                f
            );
        }
    });

    // 启动压测工作线程
    let mut handles = vec![];
    for _ in 0..1000 {
        let resolver = resolver.clone();
        let domains = domains.clone();
        let success = success.clone();
        let failure = failure.clone();

        handles.push(tokio::spawn(async move {
            let mut rng = StdRng::from_entropy();
            loop {
                let domain = domains[rng.gen_range(0..domains.len())];
                match resolver.lookup_ip(domain).await {
                    Ok(_) => success.fetch_add(1, Ordering::Relaxed),
                    Err(_) => failure.fetch_add(1, Ordering::Relaxed),
                };
            }
        }));
    }

    // 等待Ctrl+C信号退出
    tokio::signal::ctrl_c().await?;
    Ok(())
}