hbbft/tests/common_subset.rs

//! Integration tests of the Asynchronous Common Subset protocol.

extern crate hbbft;
#[macro_use]
extern crate log;
extern crate env_logger;

use std::collections::{BTreeMap, HashSet, HashMap, VecDeque};

use hbbft::common_subset;
use hbbft::common_subset::CommonSubset;
use hbbft::messaging::TargetedMessage;

type ProposedValue = Vec<u8>;

#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
struct NodeUid(usize);

/// The queue of messages of a particular Common Subset instance received by a node.
type InputQueue = VecDeque<common_subset::Message<NodeUid>>;
/// The queue of messages output from a Common Subset instance.
type OutputQueue = VecDeque<TargetedMessage<common_subset::Message<NodeUid>, NodeUid>>;

struct TestNode {
    /// Sender ID.
    id: NodeUid,
    /// The Common Subset algorithm.
    cs: CommonSubset<NodeUid>,
    /// Queue of tuples of a sender ID and a message.
    queue: VecDeque<(NodeUid, common_subset::Message<NodeUid>)>,
    /// The output of the Common Subset algorithm, if there is one.
    decision: Option<HashMap<NodeUid, ProposedValue>>,
}

impl TestNode {
    fn new(id: NodeUid, cs: CommonSubset<NodeUid>) -> TestNode {
        TestNode {
            id,
            cs,
            queue: VecDeque::new(),
            decision: None,
        }
    }

    fn handle_message(&mut self) -> (Option<HashMap<NodeUid, Vec<u8>>>, OutputQueue) {
        let (sender_id, message) = self.queue
            .pop_front()
            .expect("popping a message off the queue");
        let (output, messages) = self.cs
            .handle_message(&sender_id, message)
            .expect("handling a Common Subset message");
        debug!("{:?} produced messages: {:?}", self.id, messages);
        if let Some(ref decision) = output {
            self.decision = Some(decision.clone());
        }
        (output, messages)
    }
}

struct TestNetwork {
    nodes: BTreeMap<NodeUid, TestNode>,
    /// The next node to handle a message in its queue.
    scheduled_node_id: NodeUid,
}

impl TestNetwork {
    fn new(all_ids: &HashSet<NodeUid>) -> TestNetwork {
        let num_nodes = all_ids.len();
        // Make a node with an Agreement instance associated with the proposer node 0.
        let make_node = |id: NodeUid| {
            let cs = CommonSubset::new(id, all_ids).expect("Node creation");
            (id, TestNode::new(id, cs))
        };
        TestNetwork {
            nodes: (0..num_nodes).map(NodeUid).map(make_node).collect(),
            scheduled_node_id: NodeUid(0),
        }
    }

    fn dispatch_messages(&mut self, sender_id: NodeUid, messages: InputQueue) {
        for message in messages {
            for (id, node) in &mut self.nodes {
                if *id != sender_id {
                    debug!(
                        "Dispatching from {:?} to {:?}: {:?}",
                        sender_id, id, message
                    );
                    node.queue.push_back((sender_id, message.clone()));
                }
            }
        }
    }

    // Gets a node for receiving a message and picks the next node with a
    // non-empty message queue in a cyclic order.
    fn pick_node(&mut self) -> NodeUid {
        let id = self.scheduled_node_id;
        // Try a node with a higher ID for fairness.
        if let Some(next_id) = self.nodes
            .iter()
            .find(|(&next_id, node)| id < next_id && !node.queue.is_empty())
            .map(|(id, _)| *id)
        {
            self.scheduled_node_id = next_id;
        } else {
            // Fall back to nodes up to the currently scheduled ID.
            self.scheduled_node_id = self.nodes
                .iter()
                .find(|(&next_id, node)| id >= next_id && !node.queue.is_empty())
                .map(|(id, _)| *id)
                .expect("no more messages in any node's queue")
        }
        debug!("Picked node {:?}", self.scheduled_node_id);
        id
    }

    fn step(&mut self) -> (NodeUid, Option<HashMap<NodeUid, ProposedValue>>) {
        let sender_id = self.pick_node();
        let (output, messages) = self.nodes.get_mut(&sender_id).unwrap().handle_message();
        let messages = messages
            .into_iter()
            .map(|TargetedMessage { target: _, message }| message)
            .collect();
        self.dispatch_messages(sender_id, messages);
        (sender_id, output)
    }

    /// Make Node 0 propose a value.
    fn send_proposed_value(&mut self, sender_id: NodeUid, value: ProposedValue) {
        let messages = self.nodes
            .get_mut(&sender_id)
            .unwrap()
            .cs
            .send_proposed_value(value)
            .expect("send proposed value");
        self.dispatch_messages(
            sender_id,
            messages
                .into_iter()
                .map(|TargetedMessage { target: _, message }| message)
                .collect(),
        );
    }
}

fn test_common_subset(mut network: TestNetwork) -> BTreeMap<NodeUid, TestNode> {
    let _ = env_logger::try_init();

    // Pick the first node with a non-empty queue.
    network.pick_node();

    while network.nodes.values().any(|node| node.decision.is_none()) {
        let (NodeUid(id), output) = network.step();
        if let Some(decision) = output {
            debug!("Node {} output {:?}", id, decision);
        }
    }
    network.nodes
}

#[test]
fn test_common_subset_4_nodes() {
    let proposed_value = Vec::from("Fake news");
    let all_ids: HashSet<NodeUid> = (0..4).map(NodeUid).collect();
    let mut network = TestNetwork::new(&all_ids);
    let expected_node_decision: HashMap<NodeUid, ProposedValue> = all_ids
        .iter()
        .map(|id| (id.clone(), proposed_value.clone()))
        .collect();

    network.send_proposed_value(NodeUid(0), proposed_value.clone());
    network.send_proposed_value(NodeUid(1), proposed_value.clone());
    network.send_proposed_value(NodeUid(2), proposed_value.clone());
    network.send_proposed_value(NodeUid(3), proposed_value.clone());

    let nodes = test_common_subset(network);

    for node in nodes.values() {
        assert_eq!(
            node.decision,
            Some(expected_node_decision.clone())
        );
    }
}