1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106

// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{
    application::storage::PeerMetadataStorage,
    peer_manager::{ConnectionRequestSender, PeerManagerRequestSender},
    protocols::network::{NewNetworkEvents, NewNetworkSender},
    testutils::test_node::{
        ApplicationNetworkHandle, ApplicationNode, InboundNetworkHandle, NodeId,
        OutboundMessageReceiver,
    },
};
use channel::message_queues::QueueStyle;
use diem_config::{
    config::NodeConfig,
    network_id::{NetworkId, PeerNetworkId},
};
use std::{collections::HashMap, hash::Hash, sync::Arc, vec::Vec};

/// A trait describing a test framework for a specific application
///
/// This is essentially an abstract implementation, to get around how rust handles traits
/// there are functions to get required variables in the implementation.
///
pub trait TestFramework<Node: ApplicationNode + Sync> {
    /// Constructor for the [`TestFramework`]
    fn new(nodes: HashMap<NodeId, Node>) -> Self;

    /// A constructor for `Node` specific to the application
    fn build_node(node_id: NodeId, config: NodeConfig, peer_network_ids: &[PeerNetworkId]) -> Node;

    /// In order to have separate tasks, we have to pull these out of the framework
    fn take_node(&mut self, node_id: NodeId) -> Node;
}

/// Setup the multiple networks built for a specific node
pub fn setup_node_networks<NetworkSender: NewNetworkSender, NetworkEvents: NewNetworkEvents>(
    network_ids: &[NetworkId],
) -> (
    Vec<ApplicationNetworkHandle<NetworkSender, NetworkEvents>>,
    HashMap<NetworkId, InboundNetworkHandle>,
    HashMap<NetworkId, OutboundMessageReceiver>,
    Arc<PeerMetadataStorage>,
) {
    let mut application_handles = Vec::new();
    let mut inbound_handles = HashMap::new();
    let mut outbound_handles = HashMap::new();

    let peer_metadata_storage = PeerMetadataStorage::new(network_ids);

    // Build each individual network
    for network_id in network_ids {
        let (application_handle, inbound_handle, outbound_handle) =
            setup_network(*network_id, peer_metadata_storage.clone());
        application_handles.push(application_handle);
        inbound_handles.insert(*network_id, inbound_handle);
        outbound_handles.insert(*network_id, outbound_handle);
    }

    (
        application_handles,
        inbound_handles,
        outbound_handles,
        peer_metadata_storage,
    )
}

/// Builds all the channels used for networking
fn setup_network<NetworkSender: NewNetworkSender, NetworkEvents: NewNetworkEvents>(
    network_id: NetworkId,
    peer_metadata_storage: Arc<PeerMetadataStorage>,
) -> (
    ApplicationNetworkHandle<NetworkSender, NetworkEvents>,
    InboundNetworkHandle,
    OutboundMessageReceiver,
) {
    let (reqs_inbound_sender, reqs_inbound_receiver) = diem_channel();
    let (reqs_outbound_sender, reqs_outbound_receiver) = diem_channel();
    let (connection_outbound_sender, _connection_outbound_receiver) = diem_channel();
    let (connection_inbound_sender, connection_inbound_receiver) =
        crate::peer_manager::conn_notifs_channel::new();
    let network_sender = NetworkSender::new(
        PeerManagerRequestSender::new(reqs_outbound_sender),
        ConnectionRequestSender::new(connection_outbound_sender),
    );
    let network_events = NetworkEvents::new(reqs_inbound_receiver, connection_inbound_receiver);

    (
        (network_id, network_sender, network_events),
        InboundNetworkHandle {
            inbound_message_sender: reqs_inbound_sender,
            connection_update_sender: connection_inbound_sender,
            peer_metadata_storage,
        },
        reqs_outbound_receiver,
    )
}

/// A generic FIFO diem channel
fn diem_channel<K: Eq + Hash + Clone, T>() -> (
    channel::diem_channel::Sender<K, T>,
    channel::diem_channel::Receiver<K, T>,
) {
    static MAX_QUEUE_SIZE: usize = 8;
    channel::diem_channel::new(QueueStyle::FIFO, MAX_QUEUE_SIZE, None)
}