aptos-network 0.2.2

Aptos networking stack
Documentation 
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

// Copyright (c) Aptos
// SPDX-License-Identifier: Apache-2.0

use crate::{
    application::{
        interface::NetworkInterface,
        storage::{LockingHashMap, PeerMetadataStorage},
        types::{PeerError, PeerState},
    },
    protocols::health_checker::HealthCheckerMsg,
    transport::ConnectionMetadata,
};
use aptos_config::network_id::{NetworkId, PeerNetworkId};
use aptos_types::PeerId;
use std::{collections::hash_map::Entry, sync::Arc};

#[derive(Clone)]
struct DummySender {}

/// Dummy network so we can test the interfaces
struct DummyNetworkInterface {
    peer_metadata_storage: Arc<PeerMetadataStorage>,
}

impl NetworkInterface<HealthCheckerMsg, DummySender> for DummyNetworkInterface {
    type AppDataKey = PeerId;
    type AppData = ();

    fn peer_metadata_storage(&self) -> &PeerMetadataStorage {
        &self.peer_metadata_storage
    }

    fn sender(&self) -> DummySender {
        DummySender {}
    }

    fn app_data(&self) -> &LockingHashMap<PeerId, ()> {
        unimplemented!()
    }
}

#[test]
fn test_interface() {
    let network_id = NetworkId::Validator;
    let peer_metadata_storage = PeerMetadataStorage::test();
    let interface = DummyNetworkInterface {
        peer_metadata_storage: peer_metadata_storage.clone(),
    };
    let peer_1 = PeerId::random();
    let peer_2 = PeerId::random();
    assert_eq!(0, interface.peers(network_id).len());
    assert_eq!(0, interface.connected_peers(network_id).len());

    // Insert 2 connections, and we should have two active peers
    let connection_1 = ConnectionMetadata::mock(peer_1);
    let connection_2 = ConnectionMetadata::mock(peer_2);
    peer_metadata_storage.insert_connection(network_id, connection_1);
    peer_metadata_storage.insert_connection(network_id, connection_2.clone());
    assert_eq!(2, interface.peers(network_id).len());
    assert_eq!(2, interface.connected_peers(network_id).len());

    // Disconnecting / disconnected are not counted in active
    update_state(
        peer_metadata_storage.clone(),
        PeerNetworkId::new(network_id, peer_1),
        PeerState::Disconnecting,
    );
    assert_eq!(2, interface.peers(network_id).len());
    assert_eq!(1, interface.connected_peers(network_id).len());

    // Removing a connection with a different connection id doesn't remove it from storage
    let different_connection_2 = ConnectionMetadata::mock(peer_2);
    peer_metadata_storage.remove_connection(network_id, &different_connection_2);
    assert_eq!(2, interface.peers(network_id).len());
    assert_eq!(1, interface.connected_peers(network_id).len());

    // Removing the same connection id removes it
    peer_metadata_storage.remove_connection(network_id, &connection_2);
    assert_eq!(1, interface.peers(network_id).len());
    assert_eq!(0, interface.connected_peers(network_id).len());
}

fn update_state(
    peer_metadata_storage: Arc<PeerMetadataStorage>,
    peer_network_id: PeerNetworkId,
    state: PeerState,
) {
    peer_metadata_storage
        .write(peer_network_id, |entry| match entry {
            Entry::Vacant(..) => Err(PeerError::NotFound),
            Entry::Occupied(inner) => {
                inner.get_mut().status = state;
                Ok(())
            }
        })
        .unwrap()
}