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
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137

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

use crate::{
    application::{
        interface::NetworkInterface,
        storage::{LockingHashMap, PeerMetadataStorage},
        types::{PeerError, PeerState},
    },
    error::NetworkError,
    protocols::{
        health_checker::{
            HealthCheckerMsg, HealthCheckerNetworkEvents, HealthCheckerNetworkSender,
        },
        network::Event,
    },
};
use async_trait::async_trait;
use diem_types::PeerId;
use futures::{stream::FusedStream, Stream};
use std::{
    collections::hash_map::Entry,
    pin::Pin,
    sync::Arc,
    task::{Context, Poll},
};

#[derive(Clone, Copy, Default, Debug, Eq, PartialEq)]
pub struct HealthCheckData {
    pub round: u64,
    pub failures: u64,
}

impl HealthCheckData {
    pub fn new(round: u64) -> Self {
        HealthCheckData { round, failures: 0 }
    }
}

/// HealthChecker's view into networking
pub struct HealthCheckNetworkInterface {
    peer_metadata_storage: Arc<PeerMetadataStorage>,
    app_data: LockingHashMap<PeerId, HealthCheckData>,
    sender: HealthCheckerNetworkSender,
    receiver: HealthCheckerNetworkEvents,
}

impl HealthCheckNetworkInterface {
    pub fn new(
        peer_metadata_storage: Arc<PeerMetadataStorage>,
        sender: HealthCheckerNetworkSender,
        receiver: HealthCheckerNetworkEvents,
    ) -> Self {
        Self {
            peer_metadata_storage,
            app_data: LockingHashMap::new(),
            sender,
            receiver,
        }
    }

    /// Disconnect a peer, and keep track of the associated state
    /// Note: This removes the peer outright for now until we add GCing, and historical state management
    pub async fn disconnect_peer(&mut self, peer_id: PeerId) -> Result<(), NetworkError> {
        // Possibly already disconnected, but try anyways
        let _ = self.update_state(peer_id, PeerState::Disconnecting);
        let result = self.sender.disconnect_peer(peer_id).await;
        if result.is_ok() {
            self.remove_app_data(&peer_id);
        }
        result
    }

    pub fn connected_peers(&self) -> Vec<PeerId> {
        self.app_data.keys()
    }

    /// Update state of peer globally
    fn update_state(&self, peer_id: PeerId, state: PeerState) -> Result<(), PeerError> {
        self.peer_metadata_storage()
            .write(peer_id, |entry| match entry {
                Entry::Vacant(..) => Err(PeerError::NotFound),
                Entry::Occupied(inner) => {
                    inner.get_mut().status = state;
                    Ok(())
                }
            })
    }
}

#[async_trait]
impl NetworkInterface for HealthCheckNetworkInterface {
    type Sender = HealthCheckerNetworkSender;
    type AppData = HealthCheckData;

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

    fn sender(&self) -> Self::Sender {
        self.sender.clone()
    }

    fn insert_app_data(&self, peer_id: PeerId, data: Self::AppData) {
        self.app_data.insert(peer_id, data)
    }

    fn remove_app_data(&self, peer_id: &PeerId) {
        self.app_data.remove(peer_id)
    }

    fn read_app_data(&self, peer_id: &PeerId) -> Option<Self::AppData> {
        self.app_data.read(peer_id)
    }

    fn write_app_data<F: FnOnce(&mut Entry<PeerId, Self::AppData>) -> Result<(), PeerError>>(
        &self,
        peer_id: PeerId,
        modifier: F,
    ) -> Result<(), PeerError> {
        self.app_data.write(peer_id, modifier)
    }
}

impl Stream for HealthCheckNetworkInterface {
    type Item = Event<HealthCheckerMsg>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        Pin::new(&mut self.get_mut().receiver).poll_next(cx)
    }
}

impl FusedStream for HealthCheckNetworkInterface {
    fn is_terminated(&self) -> bool {
        self.receiver.is_terminated()
    }
}