dht/dht-search/src/verification.rs

// 负责按需调度种子可用性验证并以有界并发持久化验证结果

use std::sync::{
    Arc,
    atomic::{AtomicU64, Ordering},
};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use std::{collections::HashSet, net::SocketAddr};

use dht_crawler::DHTServer;
use dht_search::{
    domain::{InfoHash, VerificationResult},
    storage::{RocksTorrentRepository, TorrentRepository, VerificationPriority},
};
use tokio::task::{JoinHandle, JoinSet};
use tokio_util::sync::CancellationToken;

use crate::config::VerificationConfig;

#[derive(Clone)]
pub(crate) struct VerificationIngress {
    repository: Arc<RocksTorrentRepository>,
    capacity: usize,
    stats: VerificationStats,
}

#[derive(Clone, Default)]
pub(crate) struct VerificationStats {
    inner: Arc<VerificationStatsInner>,
}

#[derive(Default)]
struct VerificationStatsInner {
    accepted: AtomicU64,
    deduplicated: AtomicU64,
    rejected_full: AtomicU64,
    started: AtomicU64,
    succeeded: AtomicU64,
    failed: AtomicU64,
    peers_discovered: AtomicU64,
    handshakes_succeeded: AtomicU64,
    queue_depth: AtomicU64,
}

#[derive(Debug, Clone, Copy)]
pub(crate) struct VerificationSnapshot {
    pub(crate) accepted: u64,
    pub(crate) deduplicated: u64,
    pub(crate) rejected_full: u64,
    pub(crate) started: u64,
    pub(crate) succeeded: u64,
    pub(crate) failed: u64,
    pub(crate) peers_discovered: u64,
    pub(crate) handshakes_succeeded: u64,
    pub(crate) queue_depth: u64,
}

pub(crate) fn start(
    repository: Arc<RocksTorrentRepository>,
    server: DHTServer,
    config: VerificationConfig,
    cancel: CancellationToken,
) -> (VerificationIngress, JoinHandle<Result<(), String>>) {
    let stats = VerificationStats::default();
    stats.inner.queue_depth.store(
        repository.verification_queue_len().unwrap_or_default() as u64,
        Ordering::Relaxed,
    );
    let ingress =
        VerificationIngress::new(repository.clone(), config.queue_capacity, stats.clone());
    let task = tokio::spawn(run(repository, server, config, stats, cancel));
    (ingress, task)
}

impl VerificationIngress {
    fn new(
        repository: Arc<RocksTorrentRepository>,
        capacity: usize,
        stats: VerificationStats,
    ) -> Self {
        Self {
            repository,
            capacity,
            stats,
        }
    }

    #[cfg(test)]
    pub(crate) fn for_test(repository: Arc<RocksTorrentRepository>, capacity: usize) -> Self {
        Self::new(repository, capacity, VerificationStats::default())
    }

    pub(crate) async fn enqueue(&self, hashes: Vec<InfoHash>, priority: VerificationPriority) {
        if hashes.is_empty() {
            return;
        }
        let repository = self.repository.clone();
        let capacity = self.capacity;
        let result = tokio::task::spawn_blocking(move || {
            let outcome =
                repository.enqueue_verification(&hashes, priority, unix_timestamp(), capacity)?;
            let queue_len = repository.verification_queue_len()?;
            Ok::<_, dht_search::storage::StorageError>((outcome, queue_len))
        })
        .await;
        match result {
            Ok(Ok((outcome, queue_len))) => {
                self.stats
                    .inner
                    .accepted
                    .fetch_add(outcome.accepted as u64, Ordering::Relaxed);
                self.stats
                    .inner
                    .deduplicated
                    .fetch_add(outcome.deduplicated as u64, Ordering::Relaxed);
                self.stats
                    .inner
                    .rejected_full
                    .fetch_add(outcome.rejected_full as u64, Ordering::Relaxed);
                self.stats
                    .inner
                    .queue_depth
                    .store(queue_len as u64, Ordering::Relaxed);
            }
            Ok(Err(error)) => tracing::error!(%error, "可用性验证任务持久化失败"),
            Err(error) => tracing::error!(%error, "可用性验证入队任务异常"),
        }
    }

    pub(crate) fn stats(&self) -> VerificationStats {
        self.stats.clone()
    }
}

impl VerificationStats {
    pub(crate) fn snapshot(&self) -> VerificationSnapshot {
        VerificationSnapshot {
            accepted: self.inner.accepted.load(Ordering::Relaxed),
            deduplicated: self.inner.deduplicated.load(Ordering::Relaxed),
            rejected_full: self.inner.rejected_full.load(Ordering::Relaxed),
            started: self.inner.started.load(Ordering::Relaxed),
            succeeded: self.inner.succeeded.load(Ordering::Relaxed),
            failed: self.inner.failed.load(Ordering::Relaxed),
            peers_discovered: self.inner.peers_discovered.load(Ordering::Relaxed),
            handshakes_succeeded: self.inner.handshakes_succeeded.load(Ordering::Relaxed),
            queue_depth: self.inner.queue_depth.load(Ordering::Relaxed),
        }
    }
}

async fn run(
    repository: Arc<RocksTorrentRepository>,
    server: DHTServer,
    config: VerificationConfig,
    stats: VerificationStats,
    cancel: CancellationToken,
) -> Result<(), String> {
    let mut active = JoinSet::new();
    let mut ticker = tokio::time::interval(Duration::from_millis(config.poll_interval_millis));
    ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
    loop {
        tokio::select! {
            _ = cancel.cancelled() => break,
            result = active.join_next(), if !active.is_empty() => {
                if let Some(result) = result {
                    result.map_err(|error| error.to_string())??;
                }
            }
            _ = ticker.tick() => {
                while active.len() < config.max_active {
                    let claim_repository = repository.clone();
                    let lease_secs = config.lease_secs;
                    let request = tokio::task::spawn_blocking(move || {
                        claim_repository.claim_verification(unix_timestamp(), lease_secs)
                    })
                    .await
                    .map_err(|error| error.to_string())?
                    .map_err(|error| error.to_string())?;
                    let Some(request) = request else { break };
                    stats.inner.started.fetch_add(1, Ordering::Relaxed);
                    active.spawn(verify_one(
                        repository.clone(),
                        server.clone(),
                        request.info_hash,
                        config.max_peer_attempts,
                        stats.clone(),
                    ));
                }
            }
        }
    }
    while let Some(result) = active.join_next().await {
        result.map_err(|error| error.to_string())??;
    }
    Ok(())
}

async fn verify_one(
    repository: Arc<RocksTorrentRepository>,
    server: DHTServer,
    info_hash: InfoHash,
    max_peer_attempts: usize,
    stats: VerificationStats,
) -> Result<(), String> {
    let peer_repository = repository.clone();
    let stored_peers = tokio::task::spawn_blocking(move || peer_repository.get(info_hash))
        .await
        .map_err(|error| error.to_string())?
        .map_err(|error| error.to_string())?
        .map(|record| record.source_peers)
        .unwrap_or_default();
    let lookup = server.lookup_peers(*info_hash.as_bytes()).await;
    let dht_peers = lookup.map(|result| result.peers).unwrap_or_default();
    let mut unique = HashSet::with_capacity(dht_peers.len() + stored_peers.len());
    let mut peers = Vec::with_capacity(dht_peers.len() + stored_peers.len());
    for peer in dht_peers.into_iter().chain(
        stored_peers
            .iter()
            .filter_map(|peer| peer.parse::<SocketAddr>().ok()),
    ) {
        if unique.insert(peer) {
            peers.push(peer);
        }
    }
    stats
        .inner
        .peers_discovered
        .fetch_add(peers.len() as u64, Ordering::Relaxed);
    let mut handshakes = JoinSet::new();
    for peer in peers.iter().copied().take(max_peer_attempts) {
        let server = server.clone();
        let hash = *info_hash.as_bytes();
        handshakes.spawn(async move { server.verify_peer_handshake(hash, peer).await });
    }
    let mut reachable = 0_u32;
    while let Some(result) = handshakes.join_next().await {
        if result.map_err(|error| error.to_string())? {
            reachable = reachable.saturating_add(1);
        }
    }
    stats
        .inner
        .handshakes_succeeded
        .fetch_add(u64::from(reachable), Ordering::Relaxed);
    if reachable > 0 {
        stats.inner.succeeded.fetch_add(1, Ordering::Relaxed);
    } else {
        stats.inner.failed.fetch_add(1, Ordering::Relaxed);
    }
    let result = VerificationResult {
        verified_at: unix_timestamp(),
        discovered_peers: peers.len().min(u32::MAX as usize) as u32,
        reachable_peers: reachable,
    };
    let queue_len = tokio::task::spawn_blocking(move || {
        repository.finish_verification(info_hash, result)?;
        repository.verification_queue_len()
    })
    .await
    .map_err(|error| error.to_string())?
    .map_err(|error| error.to_string())?;
    stats
        .inner
        .queue_depth
        .store(queue_len as u64, Ordering::Relaxed);
    Ok(())
}

fn unix_timestamp() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_secs()
}