frid/crates/music-dht/src/node.rs

//! The DHT node: event handling, peer exchange, iterative lookups and
//! publication.

use std::collections::{HashMap, HashSet};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex, MutexGuard, PoisonError};
use std::time::{Duration, Instant};

use federation_net::{EndpointId, NetworkEngine, NetworkEvent, NetworkEventReceiver, PeerTicket};
use futures::future::join_all;
use futures::stream::{FuturesUnordered, StreamExt};
use tokio::sync::mpsc;
use tokio::time::timeout;
use tracing::{debug, info, warn};

use crate::config::MusicDhtConfig;
use crate::database::MusicDhtStorage;
use crate::dht::validate_store;
use crate::error::{MusicDhtError, Result};
use crate::message::{
    DHT_PROTOCOL_VERSION, FindNodeRequest, FindNodeResponse, FindValueRequest, FindValueResponse,
    Hello, MAX_PEER_EXCHANGE_CONTACTS, MAX_RECORDS_PER_BATCH, MusicDhtMessage, PeerExchange,
    PingRequest, PongResponse, RequestEnvelope, RequestId, ResponseEnvelope, StoreBatchRequest,
    StoreBatchResponse, StoreRecordRequest,
};
use crate::record::{ACTIVE_RECORD_TTL, DhtKey, LibraryItem, StoredRecord, TOMBSTONE_TTL, now_ms};
use crate::request::{DhtResponse, PendingRequests};
use crate::routing::{ALPHA, K, MAX_LOOKUP_REQUESTS, NodeContact, NodeId, RoutingTable, distance};
use crate::service::{MusicDhtEvent, PublishStats};

fn lock<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
    mutex.lock().unwrap_or_else(PoisonError::into_inner)
}

/// Base of the exponential backoff applied to a contact after a failed dial.
const DIAL_BACKOFF_BASE: Duration = Duration::from_secs(30);
/// Upper bound of the dial backoff.
const DIAL_BACKOFF_MAX: Duration = Duration::from_secs(10 * 60);
/// Consecutive failed dials after which a contact is evicted entirely.
const DIAL_FAILURES_BEFORE_EVICT: u32 = 5;

/// Dial-failure state of one currently unreachable contact.
#[derive(Debug, Clone, Copy)]
struct DialFailure {
    consecutive: u32,
    last_attempt_ms: u64,
}

/// How long a contact is skipped after `consecutive` failed dials:
/// 30s, 1m, 2m, 4m, 8m, then capped at 10 minutes.
fn dial_backoff(consecutive: u32) -> Duration {
    let exponent = consecutive.saturating_sub(1).min(8);
    DIAL_BACKOFF_BASE
        .saturating_mul(1u32 << exponent)
        .min(DIAL_BACKOFF_MAX)
}

/// An outbound DHT request, before it is wrapped in an envelope.
enum OutboundRequest {
    Ping,
    FindNode(FindNodeRequest),
    FindValue(FindValueRequest),
    StoreBatch(Box<StoreBatchRequest>),
}

/// Upper bound on the estimated payload bytes of one [`StoreBatchRequest`].
/// Kept far below federation-net's frame limit so envelope overhead and
/// estimation error never push a frame over it.
const MAX_BATCH_BYTES: usize = 128 * 1024;

/// Items processed per publish wave; bounds the memory used for per-peer
/// batch construction during a full-library republish.
const PUBLISH_WAVE_ITEMS: usize = 256;

/// Result of one iterative lookup.
pub(crate) struct LookupOutcome {
    /// Records found (value lookups only).
    pub records: Vec<StoredRecord>,
    /// Number of distinct peers actually queried.
    pub queried: usize,
    /// Number of distinct nodes known to the lookup (seeds + discovered).
    pub discovered: usize,
}

/// Shared state of one DHT node.
pub(crate) struct Node {
    pub engine: NetworkEngine<MusicDhtMessage>,
    pub db: Arc<dyn MusicDhtStorage>,
    pub config: MusicDhtConfig,
    pub node_id: NodeId,
    pub endpoint_id: EndpointId,
    routing: Mutex<RoutingTable>,
    pending: PendingRequests,
    /// Peers we already introduced ourselves to (per connection).
    hello_sent: Mutex<HashSet<EndpointId>>,
    /// Peers we already gossiped contacts to (per connection).
    exchange_sent: Mutex<HashSet<EndpointId>>,
    /// Contacts that recently failed to dial, with their backoff state.
    dial_failures: Mutex<HashMap<EndpointId, DialFailure>>,
    events: Mutex<Option<mpsc::Sender<MusicDhtEvent>>>,
    /// Set once the post-startup republish has been triggered.
    initial_republish_done: AtomicBool,
    shutting_down: AtomicBool,
}

impl Node {
    pub fn new(
        engine: NetworkEngine<MusicDhtMessage>,
        db: Arc<dyn MusicDhtStorage>,
        config: MusicDhtConfig,
        events: mpsc::Sender<MusicDhtEvent>,
    ) -> Self {
        let endpoint_id = engine.endpoint_id();
        let node_id = NodeId::from_endpoint(&endpoint_id);
        Self {
            engine,
            db,
            config,
            node_id,
            endpoint_id,
            routing: Mutex::new(RoutingTable::new(node_id)),
            pending: PendingRequests::default(),
            hello_sent: Mutex::new(HashSet::new()),
            exchange_sent: Mutex::new(HashSet::new()),
            dial_failures: Mutex::new(HashMap::new()),
            events: Mutex::new(Some(events)),
            initial_republish_done: AtomicBool::new(false),
            shutting_down: AtomicBool::new(false),
        }
    }

    pub fn is_shutting_down(&self) -> bool {
        self.shutting_down.load(Ordering::SeqCst)
    }

    pub fn begin_shutdown(&self) {
        self.shutting_down.store(true, Ordering::SeqCst);
        *lock(&self.events) = None;
    }

    pub fn ensure_running(&self) -> Result<()> {
        if self.is_shutting_down() {
            Err(MusicDhtError::ShuttingDown)
        } else {
            Ok(())
        }
    }

    async fn emit(&self, event: MusicDhtEvent) {
        let sender = lock(&self.events).clone();
        if let Some(sender) = sender {
            let _ = sender.send(event).await;
        }
    }

    /// All known DHT contacts.
    pub fn known_contacts(&self) -> Vec<NodeContact> {
        lock(&self.routing).contacts()
    }

    /// Seeds the routing table (used at startup with persisted contacts).
    pub fn seed_contacts(&self, contacts: Vec<NodeContact>) {
        let mut routing = lock(&self.routing);
        for contact in contacts {
            if contact.peer_id != self.endpoint_id {
                routing.upsert(contact);
            }
        }
    }

    /// Adds or refreshes a contact learned from the network.
    ///
    /// The node id is always re-derived from the endpoint id instead of
    /// trusting the gossiped value. The first contact ever learned triggers
    /// the post-startup republish.
    async fn upsert_contact(self: &Arc<Self>, peer_id: EndpointId, ticket: String) {
        if peer_id == self.endpoint_id {
            return;
        }
        let contact = NodeContact {
            node_id: NodeId::from_endpoint(&peer_id),
            peer_id,
            ticket,
            last_seen_ms: now_ms(),
        };
        let is_new = lock(&self.routing).upsert(contact.clone());
        if let Err(err) = self.db.upsert_known_peer(&contact).await {
            warn!(error = %err, "failed to persist known peer");
        }
        if is_new {
            info!(peer = %contact.peer_id, node = %contact.node_id, "learned new DHT contact");
            self.emit(MusicDhtEvent::ContactDiscovered {
                contact: contact.clone(),
            })
            .await;
        }
        self.maybe_trigger_initial_republish();
    }

    /// Spawns the post-startup republish once at least one contact is known.
    pub fn maybe_trigger_initial_republish(self: &Arc<Self>) {
        if lock(&self.routing).is_empty() || self.is_shutting_down() {
            return;
        }
        if self.initial_republish_done.swap(true, Ordering::SeqCst) {
            return;
        }
        let node = self.clone();
        tokio::spawn(async move {
            match node.republish_all().await {
                Ok(stats) => info!(
                    records = stats.records,
                    keys = stats.keys,
                    nodes = stats.remote_nodes,
                    "post-startup republish finished"
                ),
                Err(err) => warn!(error = %err, "post-startup republish failed"),
            }
        });
    }

    /// Consumes `federation-net` events until the engine shuts down.
    pub async fn run_event_loop(
        self: Arc<Self>,
        mut receiver: NetworkEventReceiver<MusicDhtMessage>,
    ) {
        while let Some(event) = receiver.recv().await {
            match event {
                NetworkEvent::PeerConnected { peer_id, .. } => {
                    debug!(peer = %peer_id, "peer connected");
                    self.clear_dial_failures(&peer_id);
                    self.emit(MusicDhtEvent::PeerConnected { peer_id }).await;
                    self.send_hello(peer_id).await;
                }
                NetworkEvent::PeerDisconnected { peer_id, .. } => {
                    debug!(peer = %peer_id, "peer disconnected");
                    lock(&self.hello_sent).remove(&peer_id);
                    lock(&self.exchange_sent).remove(&peer_id);
                    self.emit(MusicDhtEvent::PeerDisconnected { peer_id }).await;
                }
                NetworkEvent::MessageReceived { peer_id, message } => {
                    self.on_message(peer_id, message).await;
                }
                NetworkEvent::ProtocolError { peer_id, error } => {
                    self.emit(MusicDhtEvent::Error {
                        message: match peer_id {
                            Some(peer) => format!("transport error with {peer}: {error}"),
                            None => format!("transport error: {error}"),
                        },
                    })
                    .await;
                }
            }
        }
        debug!("network event loop finished");
    }

    async fn send_message(&self, peer: EndpointId, message: &MusicDhtMessage) -> Result<()> {
        self.engine.send(peer, message).await.map_err(Into::into)
    }

    async fn send_hello(self: &Arc<Self>, peer: EndpointId) {
        // Mark before sending so a crossing Hello does not trigger an echo.
        if !lock(&self.hello_sent).insert(peer) {
            return;
        }
        let ticket = match self.engine.ticket().await {
            Ok(ticket) => ticket.to_string(),
            Err(err) => {
                warn!(error = %err, "cannot create own ticket for hello");
                lock(&self.hello_sent).remove(&peer);
                return;
            }
        };
        let hello = MusicDhtMessage::Hello(Hello {
            node_id: self.node_id,
            peer_id: self.endpoint_id,
            ticket,
            protocol_version: DHT_PROTOCOL_VERSION,
        });
        if let Err(err) = self.send_message(peer, &hello).await {
            debug!(peer = %peer, error = %err, "failed to send hello");
            lock(&self.hello_sent).remove(&peer);
        }
    }

    async fn send_peer_exchange(self: &Arc<Self>, peer: EndpointId) {
        if !lock(&self.exchange_sent).insert(peer) {
            return;
        }
        let mut peers: Vec<NodeContact> = self
            .known_contacts()
            .into_iter()
            .filter(|contact| contact.peer_id != peer && contact.peer_id != self.endpoint_id)
            .collect();
        // Prefer the most recently seen contacts.
        peers.sort_by_key(|contact| std::cmp::Reverse(contact.last_seen_ms));
        peers.truncate(MAX_PEER_EXCHANGE_CONTACTS);
        if peers.is_empty() {
            return;
        }
        debug!(peer = %peer, count = peers.len(), "sending peer exchange");
        let message = MusicDhtMessage::PeerExchange(PeerExchange { peers });
        if let Err(err) = self.send_message(peer, &message).await {
            debug!(peer = %peer, error = %err, "failed to send peer exchange");
        }
    }

    async fn on_message(self: &Arc<Self>, peer: EndpointId, message: MusicDhtMessage) {
        lock(&self.routing).touch(&peer, now_ms());
        match message {
            MusicDhtMessage::Hello(hello) => self.on_hello(peer, hello).await,
            MusicDhtMessage::PeerExchange(exchange) => {
                self.on_peer_exchange(peer, exchange).await;
            }
            MusicDhtMessage::Ping(env) => {
                let response = MusicDhtMessage::Pong(ResponseEnvelope {
                    request_id: env.request_id,
                    payload: PongResponse {
                        node_id: self.node_id,
                    },
                });
                let _ = self.send_message(peer, &response).await;
            }
            MusicDhtMessage::FindNode(env) => {
                let nodes = self.closest_for_response(env.payload.target.as_bytes(), &peer);
                let response = MusicDhtMessage::FindNodeResult(ResponseEnvelope {
                    request_id: env.request_id,
                    payload: FindNodeResponse { nodes },
                });
                let _ = self.send_message(peer, &response).await;
            }
            MusicDhtMessage::FindValue(env) => {
                let payload = self.answer_find_value(&env.payload, &peer).await;
                let response = MusicDhtMessage::FindValueResult(ResponseEnvelope {
                    request_id: env.request_id,
                    payload,
                });
                let _ = self.send_message(peer, &response).await;
            }
            MusicDhtMessage::StoreBatch(env) => {
                let payload = self.answer_store_batch(env.payload, &peer).await;
                let response = MusicDhtMessage::StoreBatchResult(ResponseEnvelope {
                    request_id: env.request_id,
                    payload,
                });
                let _ = self.send_message(peer, &response).await;
            }
            MusicDhtMessage::Pong(env) => {
                self.pending
                    .complete(&env.request_id, &peer, DhtResponse::Pong(env.payload));
            }
            MusicDhtMessage::FindNodeResult(env) => {
                self.pending
                    .complete(&env.request_id, &peer, DhtResponse::FindNode(env.payload));
            }
            MusicDhtMessage::FindValueResult(env) => {
                self.pending
                    .complete(&env.request_id, &peer, DhtResponse::FindValue(env.payload));
            }
            MusicDhtMessage::StoreBatchResult(env) => {
                self.pending
                    .complete(&env.request_id, &peer, DhtResponse::StoreBatch(env.payload));
            }
        }
    }

    async fn on_hello(self: &Arc<Self>, peer: EndpointId, hello: Hello) {
        if hello.protocol_version != DHT_PROTOCOL_VERSION {
            warn!(peer = %peer, version = hello.protocol_version, "unsupported DHT protocol version");
            return;
        }
        // The authenticated identity comes from the connection; the id fields
        // inside the payload must be consistent with it.
        if hello.peer_id != peer || hello.node_id != NodeId::from_endpoint(&peer) {
            warn!(peer = %peer, "hello with inconsistent identity; ignoring");
            self.emit(MusicDhtEvent::Error {
                message: format!("peer {peer} sent a hello with a mismatched identity"),
            })
            .await;
            return;
        }
        debug!(peer = %peer, "received hello");
        self.upsert_contact(peer, hello.ticket).await;
        // Introduce ourselves if the remote connected first, then gossip.
        self.send_hello(peer).await;
        self.send_peer_exchange(peer).await;
    }

    async fn on_peer_exchange(self: &Arc<Self>, peer: EndpointId, exchange: PeerExchange) {
        let contacts = sanitize_peer_exchange(self.endpoint_id, peer, exchange.peers);
        let accepted = contacts.len();
        for contact in contacts {
            self.upsert_contact(contact.peer_id, contact.ticket).await;
        }
        debug!(peer = %peer, accepted, "processed peer exchange");
    }

    /// Contacts for a FindNode/FindValue response: closest to the target,
    /// excluding the requester itself.
    fn closest_for_response(&self, target: &[u8; 32], requester: &EndpointId) -> Vec<NodeContact> {
        lock(&self.routing)
            .closest(target, K + 1)
            .into_iter()
            .filter(|contact| &contact.peer_id != requester)
            .take(K)
            .collect()
    }

    async fn answer_find_value(
        &self,
        request: &FindValueRequest,
        requester: &EndpointId,
    ) -> FindValueResponse {
        match self.db.dht_records_by_key(request.key, now_ms()).await {
            Ok(records) if !records.is_empty() => FindValueResponse::Records { records },
            Ok(_) => FindValueResponse::CloserNodes {
                nodes: self.closest_for_response(request.key.as_bytes(), requester),
            },
            Err(err) => {
                warn!(error = %err, "find-value lookup in the local store failed");
                FindValueResponse::CloserNodes {
                    nodes: self.closest_for_response(request.key.as_bytes(), requester),
                }
            }
        }
    }

    /// Validates every entry of a batch individually — one bad entry never
    /// poisons the rest — and applies the valid ones in a single write
    /// transaction.
    async fn answer_store_batch(
        &self,
        request: StoreBatchRequest,
        sender: &EndpointId,
    ) -> StoreBatchResponse {
        if self.is_shutting_down() {
            return StoreBatchResponse { stored: 0 };
        }
        if request.entries.len() > MAX_RECORDS_PER_BATCH {
            warn!(
                from = %sender,
                count = request.entries.len(),
                "rejected oversized store batch"
            );
            return StoreBatchResponse { stored: 0 };
        }
        let received = request.entries.len();
        let now = now_ms();
        let mut valid = Vec::with_capacity(received);
        for entry in request.entries {
            let key = entry.key;
            match validate_store(entry, &self.config.network_id, now) {
                Ok(record) => valid.push((key, record)),
                Err(reason) => {
                    warn!(from = %sender, reason = %reason, "rejected DHT store entry");
                }
            }
        }
        let stored = match self.db.store_dht_records(valid).await {
            Ok(outcomes) => outcomes.into_iter().filter(|stored| *stored).count() as u32,
            Err(err) => {
                warn!(error = %err, "failed to store DHT record batch");
                0
            }
        };
        debug!(from = %sender, received, stored, "stored DHT record batch");
        StoreBatchResponse { stored }
    }

    /// Makes sure a connection to the contact exists, dialing its ticket if
    /// necessary. The Hello exchange runs asynchronously via the event loop.
    ///
    /// On-demand dials are bounded by the (short) `dial_timeout` and feed the
    /// dial-failure backoff, so unreachable contacts cannot stall lookups.
    async fn ensure_connected(&self, contact: &NodeContact) -> Result<EndpointId> {
        self.ensure_running()?;
        if self.engine.is_connected(contact.peer_id) {
            return Ok(contact.peer_id);
        }
        let ticket: PeerTicket = contact
            .ticket
            .parse()
            .map_err(|err| MusicDhtError::InvalidTicket(format!("{err}")))?;
        debug!(peer = %contact.peer_id, "connecting on demand");
        let result = timeout(self.config.dial_timeout, self.engine.connect(ticket))
            .await
            .map_err(|_| MusicDhtError::Timeout)
            .and_then(|res| res.map_err(Into::into));
        match result {
            Ok(peer) => {
                self.clear_dial_failures(&peer);
                Ok(peer)
            }
            Err(err) => {
                self.note_dial_failure(contact).await;
                Err(err)
            }
        }
    }

    /// Forgets the dial-failure history of a peer (it proved reachable).
    fn clear_dial_failures(&self, peer: &EndpointId) {
        lock(&self.dial_failures).remove(peer);
    }

    /// `true` if the contact recently failed to dial and its backoff window
    /// has not elapsed yet. Connected peers are never considered backed off.
    fn dial_backoff_active(&self, peer: &EndpointId, now_ms: u64) -> bool {
        if self.engine.is_connected(*peer) {
            return false;
        }
        match lock(&self.dial_failures).get(peer) {
            Some(failure) => {
                let backoff = dial_backoff(failure.consecutive).as_millis() as u64;
                now_ms < failure.last_attempt_ms.saturating_add(backoff)
            }
            None => false,
        }
    }

    /// Records a failed dial; after [`DIAL_FAILURES_BEFORE_EVICT`] failures
    /// in a row the contact is dropped from the routing table and the
    /// database (gossip re-adds it with a clean slate if it comes back).
    async fn note_dial_failure(&self, contact: &NodeContact) {
        let consecutive = {
            let mut failures = lock(&self.dial_failures);
            let failure = failures.entry(contact.peer_id).or_insert(DialFailure {
                consecutive: 0,
                last_attempt_ms: 0,
            });
            failure.consecutive += 1;
            failure.last_attempt_ms = now_ms();
            failure.consecutive
        };
        if consecutive < DIAL_FAILURES_BEFORE_EVICT {
            debug!(
                peer = %contact.peer_id,
                consecutive,
                backoff_s = dial_backoff(consecutive).as_secs(),
                "dial failed; backing off"
            );
            return;
        }
        lock(&self.dial_failures).remove(&contact.peer_id);
        lock(&self.routing).remove(&contact.peer_id);
        if let Err(err) = self.db.delete_known_peer(contact.peer_id).await {
            warn!(error = %err, "failed to delete evicted peer from the database");
        }
        info!(peer = %contact.peer_id, "evicted unreachable DHT contact");
    }

    /// Sends one request and awaits its response, cleaning up the pending
    /// entry on timeout.
    async fn request(
        &self,
        contact: &NodeContact,
        request: OutboundRequest,
    ) -> Result<DhtResponse> {
        let peer = self.ensure_connected(contact).await?;
        let request_id = RequestId::random();
        let receiver = self.pending.register(request_id, peer)?;
        // Lookups may cancel this future (early exit); the guard makes sure
        // the pending entry never outlives it.
        let _cleanup = self.pending.remove_on_drop(request_id);
        tracing::trace!(pending = self.pending.len(), peer = %peer, "sending DHT request");
        let message = match request {
            OutboundRequest::Ping => MusicDhtMessage::Ping(RequestEnvelope {
                request_id,
                payload: PingRequest {},
            }),
            OutboundRequest::FindNode(payload) => MusicDhtMessage::FindNode(RequestEnvelope {
                request_id,
                payload,
            }),
            OutboundRequest::FindValue(payload) => MusicDhtMessage::FindValue(RequestEnvelope {
                request_id,
                payload,
            }),
            OutboundRequest::StoreBatch(payload) => MusicDhtMessage::StoreBatch(RequestEnvelope {
                request_id,
                payload: *payload,
            }),
        };
        self.send_message(peer, &message).await?;
        match timeout(self.config.request_timeout, receiver).await {
            Ok(Ok(response)) => {
                lock(&self.routing).touch(&peer, now_ms());
                Ok(response)
            }
            Ok(Err(_)) => Err(MusicDhtError::Protocol("response channel closed".into())),
            Err(_) => Err(MusicDhtError::Timeout),
        }
    }

    /// Measures the round-trip time to a known contact and verifies its
    /// DHT identity.
    pub async fn ping(&self, contact: &NodeContact) -> Result<std::time::Duration> {
        let started = Instant::now();
        match self.request(contact, OutboundRequest::Ping).await? {
            DhtResponse::Pong(pong) => {
                if pong.node_id != NodeId::from_endpoint(&contact.peer_id) {
                    return Err(MusicDhtError::Protocol(
                        "pong with a mismatched node id".into(),
                    ));
                }
                lock(&self.routing).touch(&contact.peer_id, now_ms());
                Ok(started.elapsed())
            }
            _ => Err(MusicDhtError::Protocol(
                "unexpected response to ping".into(),
            )),
        }
    }

    /// Iterative Kademlia-style lookup.
    ///
    /// With `find_value: None` this is a node lookup converging on the
    /// closest known nodes to `target`; with `Some(key)` it sends `FindValue`
    /// and stops as soon as records arrive. How it stops depends on
    /// `drain_round`: `false` returns on the **first** records and cancels
    /// the in-flight requests (right when any single replica suffices, e.g.
    /// a content-id resolution); `true` awaits the rest of the started round
    /// first, so up to [`ALPHA`] peers contribute records (right for name
    /// searches, where different peers store different records under the
    /// same token key). Contacts in dial backoff are skipped. Never
    /// broadcasts: at most [`ALPHA`] requests run concurrently and at most
    /// [`MAX_LOOKUP_REQUESTS`] are sent in total, all hard-bounded by the
    /// lookup timeout.
    pub async fn lookup(
        &self,
        target: [u8; 32],
        find_value: Option<DhtKey>,
        drain_round: bool,
    ) -> LookupOutcome {
        let started = Instant::now();
        let deadline = tokio::time::Instant::now() + self.config.lookup_timeout;
        let mut candidates: Vec<NodeContact> = lock(&self.routing).closest(&target, K);
        let mut known: HashSet<EndpointId> =
            candidates.iter().map(|contact| contact.peer_id).collect();
        let mut queried: HashSet<EndpointId> = HashSet::new();
        let mut records: HashMap<(crate::record::ItemId, EndpointId), StoredRecord> =
            HashMap::new();
        let mut sent = 0usize;

        debug!(target = %NodeId::from_bytes(target), seeds = candidates.len(), "lookup started");

        'rounds: loop {
            if tokio::time::Instant::now() >= deadline {
                debug!("lookup deadline reached");
                break;
            }
            candidates.sort_by_key(|contact| distance(contact.node_id.as_bytes(), &target));
            let budget = MAX_LOOKUP_REQUESTS.saturating_sub(sent);
            let round_now = now_ms();
            let batch: Vec<NodeContact> = candidates
                .iter()
                .filter(|contact| {
                    !queried.contains(&contact.peer_id)
                        && !self.dial_backoff_active(&contact.peer_id, round_now)
                })
                .take(ALPHA.min(budget))
                .cloned()
                .collect();
            if batch.is_empty() {
                break;
            }
            sent += batch.len();
            for contact in &batch {
                queried.insert(contact.peer_id);
            }

            // Process responses as they complete: one dead contact must not
            // hold back the answers of the live ones.
            let mut in_flight: FuturesUnordered<_> = batch
                .iter()
                .map(|contact| async move {
                    let request = match find_value {
                        Some(key) => OutboundRequest::FindValue(FindValueRequest { key }),
                        None => OutboundRequest::FindNode(FindNodeRequest {
                            target: NodeId::from_bytes(target),
                        }),
                    };
                    (contact, self.request(contact, request).await)
                })
                .collect();
            while let Ok(next) = tokio::time::timeout_at(deadline, in_flight.next()).await {
                let Some((contact, result)) = next else {
                    break; // The round is complete.
                };
                let mut found_records = false;
                let nodes = match result {
                    Ok(DhtResponse::FindNode(response)) => response.nodes,
                    Ok(DhtResponse::FindValue(FindValueResponse::Records { records: found })) => {
                        for record in found {
                            let key = (record.item.id, record.item.owner);
                            match records.get(&key) {
                                Some(existing) if !record_supersedes(&record, existing) => {}
                                _ => {
                                    records.insert(key, record);
                                }
                            }
                        }
                        found_records = true;
                        Vec::new()
                    }
                    Ok(DhtResponse::FindValue(FindValueResponse::CloserNodes { nodes })) => nodes,
                    Ok(_) => Vec::new(),
                    Err(err) => {
                        debug!(peer = %contact.peer_id, error = %err, "lookup request failed");
                        Vec::new()
                    }
                };
                for node in nodes.into_iter().take(K) {
                    if node.peer_id == self.endpoint_id || !known.insert(node.peer_id) {
                        continue;
                    }
                    // Re-derive the node id instead of trusting gossip.
                    candidates.push(NodeContact {
                        node_id: NodeId::from_endpoint(&node.peer_id),
                        peer_id: node.peer_id,
                        ticket: node.ticket,
                        last_seen_ms: now_ms(),
                    });
                }
                if find_value.is_some() && found_records && !drain_round {
                    // Dropping `in_flight` cancels the outstanding requests.
                    break 'rounds;
                }
            }
            if find_value.is_some() && !records.is_empty() {
                // Drain mode: the round that produced records has completed;
                // no further rounds are needed.
                break;
            }
            if tokio::time::Instant::now() >= deadline {
                debug!("lookup deadline reached");
                break;
            }
            if sent >= MAX_LOOKUP_REQUESTS {
                debug!("lookup request budget exhausted");
                break;
            }
        }

        debug!(
            queried = queried.len(),
            discovered = known.len(),
            records = records.len(),
            elapsed_ms = started.elapsed().as_millis() as u64,
            "lookup finished"
        );
        LookupOutcome {
            records: records.into_values().collect(),
            queried: queried.len(),
            discovered: known.len(),
        }
    }

    /// Publishes item records (active or tombstones) under all their DHT
    /// keys to the closest known nodes.
    ///
    /// Publish targets come straight from the routing table instead of one
    /// iterative network lookup per key: hello/gossip and search lookups keep
    /// the table fresh, and while the network is small every peer is in it,
    /// so the K closest to any key are exact. On larger networks this is the
    /// usual Kademlia approximation — a value lookup still converges onto the
    /// same neighborhood through `CloserNodes` hops, and every record also
    /// lives on its owner. This turns a full-library republish into
    /// O(peers) batched requests instead of O(items × keys) lookups.
    pub async fn publish_items(&self, items: &[LibraryItem]) -> Result<PublishStats> {
        self.ensure_running()?;
        let mut total = PublishStats::default();
        for wave in items.chunks(PUBLISH_WAVE_ITEMS) {
            let stats = self.publish_wave(wave).await?;
            total.records += stats.records;
            total.keys += stats.keys;
            // remote_nodes counts unique nodes per wave; report the widest
            // replication seen across waves.
            total.remote_nodes = total.remote_nodes.max(stats.remote_nodes);
            total.local_replica |= stats.local_replica;
        }
        Ok(total)
    }

    /// Publishes one bounded wave of items: groups all (key, record) pairs
    /// by receiving peer and sends them as [`StoreBatchRequest`]s, one peer
    /// pipeline at a time per peer, all peers concurrently.
    async fn publish_wave(&self, items: &[LibraryItem]) -> Result<PublishStats> {
        let now = now_ms();
        let contacts: Vec<NodeContact> = self
            .known_contacts()
            .into_iter()
            .filter(|contact| !self.dial_backoff_active(&contact.peer_id, now))
            .collect();

        let mut keys_total = 0usize;
        let mut local_replica = false;
        // Batches under construction, keyed by the index into `contacts`.
        let mut per_peer: HashMap<usize, Vec<StoreRecordRequest>> = HashMap::new();
        // Own replicas, applied in one write transaction at the end.
        let mut own_replicas: Vec<(DhtKey, StoredRecord)> = Vec::new();
        // K-closest target sets are memoized per key: token keys repeat
        // heavily across the items of one artist or release.
        let mut targets_memo: HashMap<DhtKey, Vec<usize>> = HashMap::new();

        for item in items {
            let ttl = if item.deleted {
                TOMBSTONE_TTL
            } else {
                ACTIVE_RECORD_TTL
            };
            let record = StoredRecord {
                item: item.clone(),
                publisher: self.endpoint_id,
                expires_at_ms: now + ttl.as_millis() as u64,
            };
            let keys = item.dht_keys(&self.config.network_id);
            keys_total += keys.len();
            for key in keys {
                let targets = targets_memo
                    .entry(key)
                    .or_insert_with(|| closest_contact_indices(&contacts, &key, K));

                // The record belongs on this node too if it is among the K
                // closest (always true while the network is smaller than K).
                let own_distance = distance(self.node_id.as_bytes(), key.as_bytes());
                let self_is_close = targets.len() < K
                    || targets.last().is_none_or(|farthest| {
                        own_distance
                            <= distance(contacts[*farthest].node_id.as_bytes(), key.as_bytes())
                    });
                if self_is_close {
                    own_replicas.push((key, record.clone()));
                }

                for index in targets.iter() {
                    per_peer.entry(*index).or_default().push(StoreRecordRequest {
                        key,
                        record: record.clone(),
                    });
                }
            }
        }

        if !own_replicas.is_empty() {
            match self.db.store_dht_records(own_replicas).await {
                Ok(_) => local_replica = true,
                Err(err) => warn!(error = %err, "failed to store own replicas"),
            }
        }

        // One concurrent pipeline per peer; batches within a pipeline are
        // sequential so a slow peer only throttles itself.
        let sends = per_peer.into_iter().map(|(index, entries)| {
            let contact = &contacts[index];
            async move {
                let mut accepted_any = false;
                for entries in chunk_store_batches(entries) {
                    let sent = entries.len();
                    let request =
                        OutboundRequest::StoreBatch(Box::new(StoreBatchRequest { entries }));
                    match self.request(contact, request).await {
                        Ok(DhtResponse::StoreBatch(StoreBatchResponse { stored })) => {
                            accepted_any |= stored > 0;
                            if (stored as usize) < sent {
                                debug!(
                                    peer = %contact.peer_id,
                                    sent,
                                    stored,
                                    "peer declined part of the store batch"
                                );
                            }
                        }
                        Ok(_) => {}
                        Err(err) => {
                            debug!(peer = %contact.peer_id, error = %err, "store batch failed");
                        }
                    }
                }
                (contact.peer_id, accepted_any)
            }
        });
        let mut remote_nodes: HashSet<EndpointId> = HashSet::new();
        for (peer, accepted) in join_all(sends).await {
            if accepted {
                remote_nodes.insert(peer);
            }
        }
        debug!(
            items = items.len(),
            keys = keys_total,
            nodes = remote_nodes.len(),
            "published item wave"
        );
        Ok(PublishStats {
            records: items.len(),
            keys: keys_total,
            remote_nodes: remote_nodes.len(),
            local_replica,
        })
    }

    /// Republishes every local record that is still alive.
    pub async fn republish_all(&self) -> Result<PublishStats> {
        self.ensure_running()?;
        let items = self.db.local_items_for_republish(now_ms()).await?;
        let total = self.publish_items(&items).await?;
        info!(
            records = total.records,
            keys = total.keys,
            "republish finished"
        );
        Ok(total)
    }

    /// Drops expired replicas from the local store.
    pub async fn sweep_expired(&self) {
        match self.db.delete_expired_records(now_ms()).await {
            Ok(0) => {}
            Ok(count) => info!(count, "removed expired DHT records"),
            Err(err) => warn!(error = %err, "failed to sweep expired records"),
        }
    }
}

/// Indices of the up-to-`count` contacts closest to `key` by XOR distance.
fn closest_contact_indices(contacts: &[NodeContact], key: &DhtKey, count: usize) -> Vec<usize> {
    let mut indices: Vec<usize> = (0..contacts.len()).collect();
    indices.sort_by_key(|index| distance(contacts[*index].node_id.as_bytes(), key.as_bytes()));
    indices.truncate(count);
    indices
}

/// Splits store entries into batches respecting both the entry-count and the
/// estimated byte limits, so a batch always fits one transport frame.
fn chunk_store_batches(entries: Vec<StoreRecordRequest>) -> Vec<Vec<StoreRecordRequest>> {
    let mut batches = Vec::new();
    let mut current: Vec<StoreRecordRequest> = Vec::new();
    let mut current_bytes = 0usize;
    for entry in entries {
        let bytes = postcard::to_stdvec(&entry)
            .map(|encoded| encoded.len())
            .unwrap_or(MAX_BATCH_BYTES);
        if !current.is_empty()
            && (current.len() >= MAX_RECORDS_PER_BATCH
                || current_bytes.saturating_add(bytes) > MAX_BATCH_BYTES)
        {
            batches.push(std::mem::take(&mut current));
            current_bytes = 0;
        }
        current_bytes += bytes;
        current.push(entry);
    }
    if !current.is_empty() {
        batches.push(current);
    }
    batches
}

/// `true` if `candidate` should replace `existing` in a search result set.
pub(crate) fn record_supersedes(candidate: &StoredRecord, existing: &StoredRecord) -> bool {
    let (c, e) = (&candidate.item, &existing.item);
    c.revision > e.revision || (c.revision == e.revision && c.deleted && !e.deleted)
}

/// Filters an incoming peer-exchange batch: drops our own contact, the
/// sender's contact and duplicate endpoint ids, and enforces the batch cap.
pub(crate) fn sanitize_peer_exchange(
    own: EndpointId,
    sender: EndpointId,
    peers: Vec<NodeContact>,
) -> Vec<NodeContact> {
    let mut seen: HashSet<EndpointId> = HashSet::new();
    peers
        .into_iter()
        .take(MAX_PEER_EXCHANGE_CONTACTS)
        .filter(|contact| {
            contact.peer_id != own && contact.peer_id != sender && seen.insert(contact.peer_id)
        })
        .collect()
}

#[cfg(test)]
mod tests {
    use super::*;

    fn test_peer(seed: u8) -> EndpointId {
        iroh::SecretKey::from_bytes(&[seed; 32]).public()
    }

    fn contact(seed: u8) -> NodeContact {
        let peer = test_peer(seed);
        NodeContact {
            node_id: NodeId::from_endpoint(&peer),
            peer_id: peer,
            ticket: format!("fnet-test-{seed}"),
            last_seen_ms: 0,
        }
    }

    fn store_entry(name: &str) -> StoreRecordRequest {
        let owner = test_peer(1);
        let normalized = crate::normalization::normalize_name(name);
        StoreRecordRequest {
            key: DhtKey::exact(&federation_net::NetworkId::from_name("test"), &normalized),
            record: StoredRecord {
                item: LibraryItem {
                    id: crate::record::ItemId::from_bytes([7u8; 32]),
                    owner,
                    kind: crate::record::ItemKind::Artist,
                    name: name.to_string(),
                    normalized_name: normalized,
                    artist_names: Vec::new(),
                    featured_artist_names: Vec::new(),
                    year: None,
                    release_type: None,
                    release_title: None,
                    track_number: None,
                    disc_number: None,
                    duration_seconds: None,
                    content_id: None,
                    revision: 1,
                    deleted: false,
                    updated_at_ms: 0,
                },
                publisher: owner,
                expires_at_ms: 1000,
            },
        }
    }

    #[test]
    fn store_batches_respect_count_and_byte_limits() {
        // Count limit: MAX_RECORDS_PER_BATCH small entries per batch.
        let entries: Vec<_> = (0..MAX_RECORDS_PER_BATCH + 1)
            .map(|_| store_entry("Massive Attack"))
            .collect();
        let batches = chunk_store_batches(entries);
        assert_eq!(batches.len(), 2);
        assert_eq!(batches[0].len(), MAX_RECORDS_PER_BATCH);
        assert_eq!(batches[1].len(), 1);

        // Byte limit: huge entries split well before the count limit.
        let huge = "x".repeat(crate::record::MAX_ITEM_NAME_BYTES);
        let entries: Vec<_> = (0..MAX_RECORDS_PER_BATCH).map(|_| store_entry(&huge)).collect();
        for batch in chunk_store_batches(entries) {
            assert!(!batch.is_empty());
            let bytes: usize = batch
                .iter()
                .map(|entry| postcard::to_stdvec(entry).expect("encodes").len())
                .sum();
            assert!(bytes <= MAX_BATCH_BYTES);
        }

        assert!(chunk_store_batches(Vec::new()).is_empty());
    }

    #[test]
    fn closest_indices_sort_by_key_distance() {
        let contacts: Vec<NodeContact> = (1..=12u8).map(contact).collect();
        let key = DhtKey::from_bytes([0x42u8; 32]);
        let indices = closest_contact_indices(&contacts, &key, K);
        assert_eq!(indices.len(), K);
        for pair in indices.windows(2) {
            assert!(
                distance(contacts[pair[0]].node_id.as_bytes(), key.as_bytes())
                    <= distance(contacts[pair[1]].node_id.as_bytes(), key.as_bytes())
            );
        }
        // Fewer contacts than requested: everyone is a target.
        assert_eq!(closest_contact_indices(&contacts[..3], &key, K).len(), 3);
    }

    #[test]
    fn dial_backoff_doubles_and_is_capped() {
        assert_eq!(dial_backoff(1), Duration::from_secs(30));
        assert_eq!(dial_backoff(2), Duration::from_secs(60));
        assert_eq!(dial_backoff(3), Duration::from_secs(120));
        assert_eq!(dial_backoff(5), Duration::from_secs(480));
        // Capped at the maximum from the 6th failure on, even for huge counts.
        assert_eq!(dial_backoff(6), DIAL_BACKOFF_MAX);
        assert_eq!(dial_backoff(u32::MAX), DIAL_BACKOFF_MAX);
    }

    #[test]
    fn peer_exchange_drops_duplicates_self_and_sender() {
        let own = test_peer(1);
        let sender = test_peer(2);
        let peers = vec![
            contact(3),
            contact(3), // duplicate
            contact(1), // ourselves
            contact(2), // the sender
            contact(4),
        ];
        let sanitized = sanitize_peer_exchange(own, sender, peers);
        let ids: Vec<EndpointId> = sanitized.iter().map(|c| c.peer_id).collect();
        assert_eq!(ids, vec![test_peer(3), test_peer(4)]);
    }

    #[test]
    fn peer_exchange_is_capped() {
        let own = test_peer(1);
        let sender = test_peer(2);
        let peers: Vec<NodeContact> = (10..10 + MAX_PEER_EXCHANGE_CONTACTS as u8 + 8)
            .map(contact)
            .collect();
        let sanitized = sanitize_peer_exchange(own, sender, peers);
        assert_eq!(sanitized.len(), MAX_PEER_EXCHANGE_CONTACTS);
    }
}