path: root/src/swarm/peer/inbound.rs

//! Dispatch of inbound messages.

use std::time::Duration;

use anyhow::{Result, anyhow, bail};
use tokio::time::Instant;
use tracing::{debug, warn};

use super::admit::admit;
use super::epoch::observe_peer_epoch;
use super::handshake::{
    discover_peers, ensure_handshake, establish_peer, peer_swarm, remember_peer, send_hello_reply,
};
use super::keepalive::{
    heal_if_diverged, heard_from, heard_heads, local_status, note_alive, note_status,
    status_received,
};
use super::lanes::SyncLanes;
use crate::locks::MutexExt;
use crate::state::AppState;
use crate::swarm::sync::SyncEvent;
use crate::transport::{self, ParsedMessage};

/// Handles `message` from `peer` in arrival order, except that a sync message is only
/// queued here, on `lanes`: everything before it (admission, liveness, the peer's epoch)
/// is done in order, applying it isn't.
pub(crate) async fn handle_message(
    state: &AppState,
    lanes: &SyncLanes,
    peer: String,
    message: transport::ParsedMessage,
) -> Result<()> {
    if peer == state.local_id() {
        debug!("Ignoring message from self");
        return Ok(());
    }
    // PING and PONG only say the peer is alive (`keepalive::note_alive`).
    if matches!(message, ParsedMessage::Ping(_) | ParsedMessage::Pong { .. }) {
        note_alive(state, &peer, &message).await;
        return Ok(());
    }
    // WHO and ME end at the identity layer (`transport::identity`). A sync message names
    // its epoch by a tag (`EpochTag`): ours, or the sender's as we know it; one naming
    // neither is of an epoch we don't sync on, and says nothing we could record.
    let stamp = match &message {
        transport::ParsedMessage::Sync(envelope) => {
            let known = match state.get_peer_state(&peer).await {
                Some(peer_state) => peer_state.lock().await.epoch,
                None => None,
            };
            let epoch = envelope
                .epoch
                .resolve(std::iter::once(state.local_epoch()).chain(known));
            if epoch.is_none() {
                debug!(peer = %peer, epoch = %envelope.epoch, "Dropping a sync message of an epoch we don't know");
            }
            epoch.map(|epoch| (epoch, envelope.epoch_seq))
        }
        _ => message.epoch_stamp(),
    };
    let Some((epoch, epoch_seq)) = stamp else {
        return Ok(());
    };

    // A reply to a HELLO from an earlier run of ours (late, or queued by the transport
    // while we were down) says nothing about now; the peer answers our current HELLO.
    if let transport::ParsedMessage::HelloReply(reply) = &message
        && reply.run_id != state.run_id
    {
        state.stats.stale_hello_reply();
        debug!(peer = %peer, "Ignoring a HELLO_REPLY to a HELLO from an earlier run");
        return Ok(());
    }

    if !admit(state, &peer, &message).await {
        return Ok(());
    }

    remember_peer(state, &peer).await;

    let run = match &message {
        transport::ParsedMessage::Hello(hello) => Some(hello.run_id),
        transport::ParsedMessage::HelloReply(reply) => Some(reply.sender_run_id),
        _ => None,
    };

    // Any inbound message proves liveness (but from a run that said farewell: its next
    // one does). When we heard from it before tells a restarted peer how long it was
    // away (in our HELLO_REPLY).
    let peer_state = state.get_or_create_peer_state(&peer).await;
    let mut hello_now = false;
    let (reactivated, previously_seen) =
        heard_from(state, &peer, &peer_state, Instant::now(), |peer_state| {
            match run {
                Some(run) => peer_state.returned(run),
                // Up and reaching us, while we wait for its HELLO_REPLY.
                None => hello_now = peer_state.hello_again_now(Instant::now()),
            }
        })
        .await;
    if hello_now {
        state.peers_due.now();
    }
    let same_epoch = observe_peer_epoch(state, &peer, epoch, epoch_seq, run).await;

    match message {
        transport::ParsedMessage::Hello(hello) => {
            // A HELLO means the peer is (re)starting its handshake with us, e.g. after a
            // restart. It becomes a member right away; our reply makes us one for it.
            let Some(swarm_id) = state.swarm_id().await else {
                // We are joining ourselves; the peer retries until we know our swarm.
                // It is up and reaches us: if we wait for its HELLO_REPLY, our HELLO or
                // the reply was lost, and ours goes again now.
                debug!(peer = %peer, "HELLO before we know our swarm; not answering yet");
                if peer_state.lock().await.hello_again_now(Instant::now()) {
                    state.peers_due.now();
                }
                return Ok(());
            };
            if let Some(theirs) = &hello.swarm_id
                && theirs != &swarm_id
            {
                bail!("Incompatible swarm at {peer}: ours is {swarm_id}, it claims {theirs}");
            }
            // A resent HELLO (our reply was lost) or a late copy of one repeats a
            // handshake we answered already: it must not reset the session again.
            let new_handshake = {
                let peer_state = state.get_or_create_peer_state(&peer).await;
                let mut peer_state = peer_state.lock().await;
                let new = peer_state.last_hello.is_none_or(|(run, handshake)| {
                    run != hello.run_id || hello.handshake > handshake
                });
                if new {
                    peer_state.last_hello = Some((hello.run_id, hello.handshake));
                }
                new
            };
            establish_peer(state, &peer, swarm_id, new_handshake, hello.ready).await?;

            let away = previously_seen.map(|seen| seen.elapsed());
            if let Err(error) = send_hello_reply(state, &peer, hello.run_id, away).await {
                // The peer retries its HELLO.
                warn!(peer = %peer, error = %format!("{error:#}"), "HELLO_REPLY failed");
            }

            let _ = state
                .sync_events
                .send(SyncEvent::PeerValidated(peer.clone()));
        }
        transport::ParsedMessage::HelloReply(reply) => {
            state.lock_rtt().hello_reply_received(&peer, Instant::now());
            crate::swarm::stability::note_away(state, reply.away_ms.map(Duration::from_millis));
            establish_peer(state, &peer, reply.swarm_id.clone(), false, reply.ready).await?;

            discover_peers(state, &peer, &reply.peers).await;

            state
                .sync_events
                .send(SyncEvent::PeerValidated(peer.clone()))
                .map_err(|error| anyhow!("Failed to notify sync worker: {error}"))?;
        }
        transport::ParsedMessage::Sync(envelope) => {
            if reactivated {
                ensure_handshake(state, &peer).await;
            }
            if !same_epoch {
                debug!(peer = %peer, epoch = %epoch, "Dropping sync message from another epoch");
                return Ok(());
            }
            crate::swarm::source::synced_from(state, &peer);
            lanes.enqueue(&peer, envelope);
        }
        transport::ParsedMessage::Blob(blob) => {
            // How long it took to come, and of a chunk, the throughput (it was asked
            // for: the request's round trip is no exchange of sync bytes to time).
            if let Some(delay_ms) = crate::swarm::hop::note(state, &peer, blob.sent).await
                && delay_ms > 0
                && let transport::BlobBody::Chunk { data, .. }
                | transport::BlobBody::DeflatedChunk { data, .. } = &blob.body
            {
                let rtt =
                    state
                        .lock_rtt()
                        .effective(&peer, Instant::now(), state.timing.rtt_fresh_for());
                state.throughput.lock_or_recover().message(
                    &peer,
                    data.len() as u64,
                    Duration::from_millis(delay_ms as u64),
                    Instant::now(),
                    rtt,
                );
            }
            crate::swarm::blob::handle(state, &peer, blob).await;
        }
        transport::ParsedMessage::Status(remote) => {
            // An announcement is no keepalive; it is answered only if we lack some of the
            // heads it announced (`relay::lacks`), so the announcer learns what we have
            // that it lacks while we pull. Having all of them, we are as far or further,
            // and what we have further we announced ourselves; differing heads alone are
            // no reason, as changes on their way to us make them differ while writes go on.
            if !remote.announcing {
                status_received(state, &peer).await;
            }
            // Announcements too: a neighbour heard from keeps none of its keepalives, so
            // while writes go on its announcements are all it says of its heads.
            heard_heads(state, &peer, &remote);
            let local = on_status(state, &peer, &remote, reactivated).await;
            if remote.announcing {
                if local.heads != remote.heads {
                    // Off the consumer: telling which heads we lack takes the document's
                    // lock.
                    let (state, peer) = (state.clone(), peer.clone());
                    tokio::spawn(async move {
                        if crate::swarm::relay::lacks(&state, &remote.heads).await
                            && let Err(error) = state
                                .transport
                                .send_message(&peer, &ParsedMessage::StatusReply(local))
                                .await
                        {
                            debug!(peer = %peer, error = %format!("{error:#}"), "STATUS_REPLY failed");
                        }
                    });
                }
                return Ok(());
            }
            if let Err(error) = state
                .transport
                .send_message(&peer, &ParsedMessage::StatusReply(local))
                .await
            {
                debug!(peer = %peer, error = %format!("{error:#}"), "STATUS_REPLY failed");
            }
        }
        transport::ParsedMessage::StatusReply(remote) => {
            // An answer to our slow round's announcement: a tree missed.
            let announced = peer_state.lock().await.slow_announced_at.take();
            if let Some(at) = announced {
                crate::swarm::relay::slow_round_answered(state, at.elapsed());
            }
            // Keepalives go one way between a pair (the one with the shorter interval
            // takes the turns): the other side hears the peer's heads in its replies.
            heard_heads(state, &peer, &remote);
            on_status(state, &peer, &remote, reactivated).await;
        }
        transport::ParsedMessage::Links(request) => {
            crate::swarm::links::answer(state, &peer, request).await;
        }
        transport::ParsedMessage::LinksReply(reply) => {
            crate::swarm::links::received(state, &peer, reply);
        }
        transport::ParsedMessage::Who(_)
        | transport::ParsedMessage::Me(_)
        | transport::ParsedMessage::Ping(_)
        | transport::ParsedMessage::Pong { .. } => {}
    }

    Ok(())
}

/// What a STATUS and a STATUS_REPLY have in common: a handshake if needed, the peer's status
/// noted, and the heal if our documents differ. Returns our status for the peer.
async fn on_status(
    state: &AppState,
    peer: &str,
    remote: &transport::Status,
    reactivated: bool,
) -> transport::Status {
    if reactivated || peer_swarm(state, peer).await.is_none() {
        ensure_handshake(state, peer).await;
    }
    note_status(state, peer, remote).await;
    let local = local_status(state, peer).await;
    heal_if_diverged(state, peer, &local, remote).await;
    local
}