path: root/src/transport/sim.rs

//! A simulated network for tests: many nodes in one process, each with a
//! [`SimTransport`] on a shared [`SimNet`].
//!
//! Messages go through the Mycelium wire encoding (`mycelium::encode_message`), so the
//! codec is exercised as in production. The net delays, loses and duplicates them per
//! link, from a seeded generator; nodes can be taken down (everything to and from them
//! is lost) and brought up again; and a test can inject a message as if a node had sent
//! it, e.g. one that arrives late.
//!
//! Time is tokio's real clock, not a paused one: a node's document work runs on the
//! blocking pool (`spawn_blocking`), which a paused clock doesn't wait for, so timers
//! would fire while the work is still running. Tests use a short keepalive instead.
//! Delays and losses come from the seed; when a message's turn comes still depends on
//! the scheduler, so a run is reproducible in its network decisions, not in its
//! interleavings.

use std::collections::{HashMap, HashSet};
use std::sync::{Arc, Mutex, MutexGuard};
use std::time::Duration;

use anyhow::{Result, bail};
use async_trait::async_trait;
use bytes::Bytes;
use tokio::sync::{mpsc, oneshot};
use tracing::{debug, warn};

use super::mycelium::{decode_message, encode_message};
use super::{EncodedMessage, MessageHandler, MessageKind, ParsedMessage, Transport};

/// How one direction of a link behaves.
#[derive(Clone, Debug)]
pub(crate) struct LinkParams {
    /// Probability that a message is lost.
    pub(crate) loss: f64,
    /// Probability that a message arrives twice (each copy with its own delay).
    pub(crate) duplicate: f64,
    /// Delay, uniform in this range. Overlapping delays reorder messages.
    pub(crate) delay: (Duration, Duration),
    /// Messages of these kinds are always lost: a peer that answers some messages but
    /// not others (e.g. keepalives but not sync).
    pub(crate) drop: Vec<MessageKind>,
}

impl Default for LinkParams {
    /// A clean, fast link: nothing lost, a few milliseconds.
    fn default() -> Self {
        Self {
            loss: 0.0,
            duplicate: 0.0,
            delay: (Duration::from_millis(1), Duration::from_millis(5)),
            drop: Vec::new(),
        }
    }
}

/// A message on its way into a node.
struct Delivery {
    from: String,
    body: Bytes,
    payload_len: usize,
    /// Told once the node has handled the message (injections wait for it).
    handled: Option<oneshot::Sender<()>>,
}

/// What the net counted, for assertions.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub(crate) struct NetCounts {
    pub(crate) sent: u64,
    pub(crate) delivered: u64,
    pub(crate) lost: u64,
    pub(crate) duplicated: u64,
}

struct NetInner {
    rng: SplitMix64,
    default_link: LinkParams,
    links: HashMap<(String, String), LinkParams>,
    down: HashSet<String>,
    inboxes: HashMap<String, mpsc::UnboundedSender<Delivery>>,
    counts: NetCounts,
}

/// The network all simulated nodes share. Cloning is cheap.
#[derive(Clone)]
pub(crate) struct SimNet {
    inner: Arc<Mutex<NetInner>>,
}

impl SimNet {
    pub(crate) fn new(seed: u64) -> Self {
        Self {
            inner: Arc::new(Mutex::new(NetInner {
                rng: SplitMix64(seed),
                default_link: LinkParams::default(),
                links: HashMap::new(),
                down: HashSet::new(),
                inboxes: HashMap::new(),
                counts: NetCounts::default(),
            })),
        }
    }

    fn lock(&self) -> MutexGuard<'_, NetInner> {
        self.inner
            .lock()
            .unwrap_or_else(|poisoned| poisoned.into_inner())
    }

    /// A transport for a new node `id` on this net.
    pub(crate) fn transport(&self, id: &str) -> Arc<SimTransport> {
        let (tx, rx) = mpsc::unbounded_channel();
        let previous = self.lock().inboxes.insert(id.to_owned(), tx);
        assert!(previous.is_none(), "node {id} is already on the net");
        Arc::new(SimTransport {
            id: id.to_owned(),
            net: self.clone(),
            inbox: tokio::sync::Mutex::new(Some(rx)),
        })
    }

    /// How links without their own parameters behave.
    pub(crate) fn set_default_link(&self, params: LinkParams) {
        self.lock().default_link = params;
    }

    /// How messages from `from` to `to` are carried (one direction).
    pub(crate) fn set_link(&self, from: &str, to: &str, params: LinkParams) {
        self.lock()
            .links
            .insert((from.to_owned(), to.to_owned()), params);
    }

    /// Everything to and from `id` is lost until [`bring_up`](Self::bring_up), including
    /// messages already on their way. The node itself keeps running.
    pub(crate) fn take_down(&self, id: &str) {
        self.lock().down.insert(id.to_owned());
    }

    pub(crate) fn bring_up(&self, id: &str) {
        self.lock().down.remove(id);
    }

    pub(crate) fn counts(&self) -> NetCounts {
        self.lock().counts
    }

    /// Hands a message to the net: decides now whether and when it arrives.
    fn send(&self, from: &str, message: &EncodedMessage) {
        let delays = {
            let mut guard = self.lock();
            let inner = &mut *guard;
            inner.counts.sent += 1;
            if inner.down.contains(from) || inner.down.contains(&message.peer) {
                inner.counts.lost += 1;
                return;
            }
            let params = inner
                .links
                .get(&(from.to_owned(), message.peer.clone()))
                .unwrap_or(&inner.default_link)
                .clone();
            if params.drop.contains(&message.kind) || inner.rng.chance(params.loss) {
                inner.counts.lost += 1;
                return;
            }
            let copies = if inner.rng.chance(params.duplicate) {
                inner.counts.duplicated += 1;
                2
            } else {
                1
            };
            (0..copies)
                .map(|_| inner.rng.duration_in(params.delay))
                .collect::<Vec<_>>()
        };
        for delay in delays {
            let net = self.clone();
            let delivery = Delivery {
                from: from.to_owned(),
                body: message.body.clone(),
                payload_len: message.payload_len,
                handled: None,
            };
            let to = message.peer.clone();
            tokio::spawn(async move {
                tokio::time::sleep(delay).await;
                net.arrive(&to, delivery);
            });
        }
    }

    /// The message reaches `to`, unless either end went down meanwhile.
    fn arrive(&self, to: &str, delivery: Delivery) {
        let mut guard = self.lock();
        let inner = &mut *guard;
        if inner.down.contains(to) || inner.down.contains(&delivery.from) {
            inner.counts.lost += 1;
            return;
        }
        // An unknown recipient (an address no node has) loses the message, as
        // Mycelium would.
        let Some(inbox) = inner.inboxes.get(to) else {
            inner.counts.lost += 1;
            return;
        };
        if inbox.send(delivery).is_ok() {
            inner.counts.delivered += 1;
        } else {
            inner.counts.lost += 1;
        }
    }

    /// Delivers `message` to `to` as if `from` had sent it, now, bypassing links and
    /// `take_down`; returns once `to` has handled it. `from` needn't be a node on the
    /// net. For messages a real node wouldn't send at that moment, e.g. one sent long
    /// ago and delivered late.
    pub(crate) async fn inject(&self, from: &str, to: &str, message: &ParsedMessage) -> Result<()> {
        let encoded = encode_message(to, message)?;
        let (handled_tx, handled_rx) = oneshot::channel();
        let delivery = Delivery {
            from: from.to_owned(),
            body: encoded.body,
            payload_len: encoded.payload_len,
            handled: Some(handled_tx),
        };
        {
            let inner = self.lock();
            let Some(inbox) = inner.inboxes.get(to) else {
                bail!("no node {to} on the net");
            };
            if inbox.send(delivery).is_err() {
                bail!("node {to} has no consumer running");
            }
        }
        handled_rx.await?;
        Ok(())
    }

    fn is_known(&self, id: &str) -> bool {
        self.lock().inboxes.contains_key(id)
    }
}

/// One node's end of the [`SimNet`]. Its identity is the ID it was created with; seeds
/// are node IDs.
pub(crate) struct SimTransport {
    id: String,
    net: SimNet,
    inbox: tokio::sync::Mutex<Option<mpsc::UnboundedReceiver<Delivery>>>,
}

#[async_trait]
impl Transport for SimTransport {
    fn local_id(&self) -> &str {
        &self.id
    }

    /// A seed is a node ID; one no node has doesn't resolve (yet).
    async fn resolve_seed(&self, seed: &str) -> Result<String> {
        if self.net.is_known(seed) {
            Ok(seed.to_owned())
        } else {
            bail!("no node {seed} on the net")
        }
    }

    fn encode(&self, peer: &str, message: &ParsedMessage) -> Result<EncodedMessage> {
        encode_message(peer, message)
    }

    /// Always accepted: the net loses messages silently, like Mycelium once it has
    /// taken them.
    async fn send_encoded(&self, message: &EncodedMessage) -> Result<()> {
        self.net.send(&self.id, message);
        Ok(())
    }

    /// Handles messages one at a time, in arrival order, as the Mycelium consumer does.
    async fn start_consumer(&self, handler: MessageHandler) -> Result<()> {
        let Some(mut inbox) = self.inbox.lock().await.take() else {
            bail!("consumer of {} already started", self.id);
        };
        while let Some(delivery) = inbox.recv().await {
            match decode_message(&delivery.body) {
                Ok(message) => {
                    debug!(node = %self.id, from = %delivery.from, message = ?message, "Simulated delivery");
                    if let Err(error) = handler(delivery.from, message, delivery.payload_len).await
                    {
                        warn!(node = %self.id, error = %format!("{error:#}"), "Failed to process inbound message");
                    }
                }
                Err(error) => {
                    warn!(node = %self.id, error = %format!("{error:#}"), "Undecodable simulated message");
                }
            }
            if let Some(handled) = delivery.handled {
                let _ = handled.send(());
            }
        }
        Ok(())
    }

    /// A simulated node has no daemon that could exit.
    async fn wait_for_exit(&self) -> std::io::Result<std::process::ExitStatus> {
        std::future::pending().await
    }
}

/// A small, seedable generator (SplitMix64): enough for network decisions, and no
/// dependency.
struct SplitMix64(u64);

impl SplitMix64 {
    fn next_u64(&mut self) -> u64 {
        self.0 = self.0.wrapping_add(0x9e37_79b9_7f4a_7c15);
        let mut z = self.0;
        z = (z ^ (z >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9);
        z = (z ^ (z >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb);
        z ^ (z >> 31)
    }

    /// Uniform in [0, 1).
    fn next_f64(&mut self) -> f64 {
        (self.next_u64() >> 11) as f64 / (1_u64 << 53) as f64
    }

    fn chance(&mut self, probability: f64) -> bool {
        probability > 0.0 && self.next_f64() < probability
    }

    fn duration_in(&mut self, (min, max): (Duration, Duration)) -> Duration {
        if max <= min {
            return min;
        }
        min + (max - min).mul_f64(self.next_f64())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::document::replica::Epoch;
    use crate::transport::{HeadsSummary, Keepalive};

    fn ping() -> ParsedMessage {
        ParsedMessage::Ping(Keepalive {
            heads: HeadsSummary::of(Vec::new()),
            epoch: Epoch::GENESIS,
            epoch_seq: 0,
            ready: true,
            syncs_with_you: false,
        })
    }

    /// Collects what a node receives.
    fn record(transport: Arc<SimTransport>) -> Arc<Mutex<Vec<String>>> {
        let received = Arc::new(Mutex::new(Vec::new()));
        let sink = Arc::clone(&received);
        tokio::spawn(async move {
            transport
                .start_consumer(Box::new(move |from, _message, _len| {
                    sink.lock().unwrap().push(from);
                    Box::pin(async { Ok(()) })
                }))
                .await
        });
        received
    }

    #[tokio::test]
    async fn delivers_decodes_and_loses_when_down() {
        let net = SimNet::new(1);
        let a = net.transport("a");
        let received = record(net.transport("b"));

        let settle = || tokio::time::sleep(Duration::from_millis(30));
        a.send_message("b", &ping()).await.unwrap();
        settle().await;
        net.take_down("b");
        a.send_message("b", &ping()).await.unwrap();
        settle().await;
        net.bring_up("b");
        a.send_message("b", &ping()).await.unwrap();
        settle().await;

        assert_eq!(*received.lock().unwrap(), vec!["a", "a"]);
        let counts = net.counts();
        assert_eq!((counts.sent, counts.delivered, counts.lost), (3, 2, 1));
        assert_eq!(a.resolve_seed("b").await.unwrap(), "b");
        assert!(a.resolve_seed("c").await.is_err());
    }

    #[tokio::test]
    async fn loss_and_duplication_follow_the_seed() {
        let run = |seed| async move {
            let net = SimNet::new(seed);
            net.set_default_link(LinkParams {
                loss: 0.3,
                duplicate: 0.2,
                ..LinkParams::default()
            });
            let a = net.transport("a");
            let _b = net.transport("b");
            for _ in 0..200 {
                a.send_message("b", &ping()).await.unwrap();
            }
            tokio::time::sleep(Duration::from_millis(50)).await;
            net.counts()
        };
        let first = run(7).await;
        assert_eq!(first, run(7).await);
        assert!(first.lost > 30 && first.lost < 90, "{first:?}");
        assert!(first.duplicated > 10, "{first:?}");
    }

    #[tokio::test]
    async fn a_link_overrides_the_default_in_one_direction() {
        let net = SimNet::new(1);
        net.set_link(
            "a",
            "b",
            LinkParams {
                loss: 1.0,
                ..LinkParams::default()
            },
        );
        let a = net.transport("a");
        let b = net.transport("b");
        let at_a = record(Arc::clone(&a));
        let at_b = record(Arc::clone(&b));

        a.send_message("b", &ping()).await.unwrap();
        b.send_message("a", &ping()).await.unwrap();
        tokio::time::sleep(Duration::from_millis(30)).await;

        assert!(at_b.lock().unwrap().is_empty());
        assert_eq!(*at_a.lock().unwrap(), vec!["b"]);
    }

    #[tokio::test]
    async fn inject_returns_once_handled() {
        let net = SimNet::new(1);
        let received = record(net.transport("b"));
        net.inject("ghost", "b", &ping()).await.unwrap();
        assert_eq!(*received.lock().unwrap(), vec!["ghost"]);
        assert!(net.inject("ghost", "nobody", &ping()).await.is_err());
    }
}