From c4f2023e2a0f19a7f851a62b9ba94400c11d28da Mon Sep 17 00:00:00 2001
From: 0xa <hit.cs.lijun@gmail.com>
Date: Fri, 22 Oct 2021 14:32:21 +0800
Subject: [PATCH] network: Add peers state metrics (#2980)

---
 network-p2p/peerset/src/lib.rs                |  7 ++++--
 network-p2p/src/metrics.rs                    |  8 ++++++
 network-p2p/src/protocol.rs                   |  6 +++--
 .../src/protocol/generic_proto/behaviour.rs   |  5 +++-
 network-p2p/src/service.rs                    | 25 ++++++++++++++++++-
 5 files changed, 45 insertions(+), 6 deletions(-)

diff --git a/network-p2p/peerset/src/lib.rs b/network-p2p/peerset/src/lib.rs
index cf5cdaed0c..e69a75d78e 100644
--- a/network-p2p/peerset/src/lib.rs
+++ b/network-p2p/peerset/src/lib.rs
@@ -32,7 +32,7 @@
 //! In addition, for each, set, the peerset also holds a list of reserved nodes towards which it
 //! will at all time try to maintain a connection with.
 
-mod peersstate;
+pub mod peersstate;
 
 use futures::prelude::*;
 use log::{debug, error, trace};
@@ -47,6 +47,7 @@ use std::{
 use std::{collections::HashSet, collections::VecDeque};
 use wasm_timer::Instant;
 
+use crate::peersstate::PeersState;
 use futures::channel::oneshot;
 use futures::channel::oneshot::{Receiver, Sender};
 pub use libp2p::PeerId;
@@ -599,7 +600,9 @@ impl Peerset {
             }
         }
     }
-
+    pub fn get_peers_state(&self) -> PeersState {
+        self.data.clone()
+    }
     /// Indicate that we received an incoming connection. Must be answered either with
     /// a corresponding `Accept` or `Reject`, except if we were already connected to this peer.
     ///
diff --git a/network-p2p/src/metrics.rs b/network-p2p/src/metrics.rs
index ce86f11f20..dbdad5b15b 100644
--- a/network-p2p/src/metrics.rs
+++ b/network-p2p/src/metrics.rs
@@ -32,6 +32,7 @@ pub struct Metrics {
     pub requests_out_failure_total: UIntCounterVec,
     pub requests_out_success_total: HistogramVec,
     pub requests_out_started_total: UIntCounterVec,
+    pub peerset_nodes: UIntGaugeVec,
 }
 
 impl Metrics {
@@ -275,6 +276,13 @@ impl Metrics {
                 )?,
                 registry,
             )?,
+            peerset_nodes: register(
+                UIntGaugeVec::new(
+                    Opts::new("sub_libp2p_peerset_nodes", "nodes numbers in each peer set"),
+                    &["sets"],
+                )?,
+                registry,
+            )?,
         })
     }
 }
diff --git a/network-p2p/src/protocol.rs b/network-p2p/src/protocol.rs
index d04de3d920..57c596f4db 100644
--- a/network-p2p/src/protocol.rs
+++ b/network-p2p/src/protocol.rs
@@ -17,7 +17,7 @@ use libp2p::swarm::{IntoProtocolsHandler, ProtocolsHandler};
 use libp2p::swarm::{NetworkBehaviour, NetworkBehaviourAction, PollParameters};
 use libp2p::PeerId;
 use log::Level;
-use sc_peerset::SetId;
+use sc_peerset::{peersstate::PeersState, SetId};
 use starcoin_types::startup_info::{ChainInfo, ChainStatus};
 use std::borrow::Cow;
 use std::collections::{HashMap, HashSet};
@@ -447,7 +447,9 @@ impl Protocol {
     pub fn peerset_debug_info(&mut self) -> serde_json::Value {
         self.behaviour.peerset_debug_info()
     }
-
+    pub fn peerset_info(&self) -> PeersState {
+        self.behaviour.peerset_info()
+    }
     pub fn on_notify(
         &mut self,
         who: PeerId,
diff --git a/network-p2p/src/protocol/generic_proto/behaviour.rs b/network-p2p/src/protocol/generic_proto/behaviour.rs
index 7df9ed862f..3cb4548973 100644
--- a/network-p2p/src/protocol/generic_proto/behaviour.rs
+++ b/network-p2p/src/protocol/generic_proto/behaviour.rs
@@ -29,6 +29,7 @@ use libp2p::swarm::{
 use log::{debug, error, trace, warn};
 use parking_lot::RwLock;
 use rand::distributions::{Distribution as _, Uniform};
+use sc_peerset::peersstate::PeersState;
 use smallvec::SmallVec;
 use std::task::{Context, Poll};
 use std::{
@@ -643,7 +644,9 @@ impl GenericProto {
     pub fn peerset_debug_info(&mut self) -> serde_json::Value {
         self.peerset.debug_info()
     }
-
+    pub fn peerset_info(&self) -> PeersState {
+        self.peerset.get_peers_state()
+    }
     /// Function that is called when the peerset wants us to connect to a peer.
     fn peerset_report_connect(&mut self, peer_id: PeerId, set_id: sc_peerset::SetId) {
         // If `PeerId` is unknown to us, insert an entry, start dialing, and return early.
diff --git a/network-p2p/src/service.rs b/network-p2p/src/service.rs
index 8dab9d9dc3..332a761e64 100644
--- a/network-p2p/src/service.rs
+++ b/network-p2p/src/service.rs
@@ -69,7 +69,7 @@ use libp2p::{kad::record, PeerId};
 use log::{error, info, trace, warn};
 use network_p2p_types::IfDisconnected;
 use parking_lot::Mutex;
-use sc_peerset::{PeersetHandle, ReputationChange};
+use sc_peerset::{peersstate, PeersetHandle, ReputationChange};
 use starcoin_metrics::{Histogram, HistogramVec};
 use starcoin_types::startup_info::ChainStatus;
 use std::collections::HashMap;
@@ -1510,6 +1510,29 @@ impl Future for NetworkWorker {
                     .connection_counters()
                     .num_pending() as u64,
             );
+            let mut peers_state = this
+                .network_service
+                .behaviour_mut()
+                .user_protocol_mut()
+                .peerset_info();
+            for set_index in 0..peers_state.num_sets() {
+                let node_count = peers_state
+                    .peers()
+                    .cloned()
+                    .collect::<Vec<_>>()
+                    .into_iter()
+                    .filter(|peer_id| {
+                        matches!(
+                            peers_state.peer(set_index, peer_id),
+                            peersstate::Peer::Connected(_)
+                        )
+                    })
+                    .count();
+                metrics
+                    .peerset_nodes
+                    .with_label_values(&[format!("set_{}", set_index).as_str()])
+                    .set(node_count as u64)
+            }
         }
 
         Poll::Pending