From dc6204601927d346fb1812bf2482457895f6b3f5 Mon Sep 17 00:00:00 2001
From: Roman Khimov <roman@nspcc.ru>
Date: Wed, 12 Oct 2022 22:29:55 +0300
Subject: [PATCH] network: add network size estimation metric

---
 pkg/network/discovery.go  |  9 +++++++++
 pkg/network/prometheus.go | 13 +++++++++++++
 2 files changed, 22 insertions(+)

diff --git a/pkg/network/discovery.go b/pkg/network/discovery.go
index e6580139d..e95d5db1f 100644
--- a/pkg/network/discovery.go
+++ b/pkg/network/discovery.go
@@ -87,6 +87,7 @@ func (d *DefaultDiscovery) BackFill(addrs ...string) {
 		d.unconnectedAddrs[addr] = connRetries
 		d.pushToPoolOrDrop(addr)
 	}
+	d.updateNetSize()
 	d.lock.Unlock()
 }
 
@@ -127,6 +128,7 @@ func (d *DefaultDiscovery) RegisterBadAddr(addr string) {
 		delete(d.unconnectedAddrs, addr)
 		delete(d.goodAddrs, addr)
 	}
+	d.updateNetSize()
 	d.lock.Unlock()
 }
 
@@ -181,6 +183,7 @@ func (d *DefaultDiscovery) RegisterGoodAddr(s string, c capability.Capabilities)
 func (d *DefaultDiscovery) UnregisterConnectedAddr(s string) {
 	d.lock.Lock()
 	delete(d.connectedAddrs, s)
+	d.updateNetSize()
 	d.lock.Unlock()
 }
 
@@ -189,9 +192,15 @@ func (d *DefaultDiscovery) RegisterConnectedAddr(addr string) {
 	d.lock.Lock()
 	delete(d.unconnectedAddrs, addr)
 	d.connectedAddrs[addr] = true
+	d.updateNetSize()
 	d.lock.Unlock()
 }
 
+// updateNetSize updates network size estimation metric. Must be called under read lock.
+func (d *DefaultDiscovery) updateNetSize() {
+	updateNetworkSizeMetric(len(d.connectedAddrs) + len(d.unconnectedAddrs))
+}
+
 func (d *DefaultDiscovery) tryAddress(addr string) {
 	err := d.transport.Dial(addr, d.dialTimeout)
 	d.lock.Lock()
diff --git a/pkg/network/prometheus.go b/pkg/network/prometheus.go
index dab35c44c..31088e3f7 100644
--- a/pkg/network/prometheus.go
+++ b/pkg/network/prometheus.go
@@ -6,6 +6,14 @@ import (
 
 // Metric used in monitoring service.
 var (
+	estimatedNetworkSize = prometheus.NewGauge(
+		prometheus.GaugeOpts{
+			Help:      "Estimated network size",
+			Name:      "network_size",
+			Namespace: "neogo",
+		},
+	)
+
 	peersConnected = prometheus.NewGauge(
 		prometheus.GaugeOpts{
 			Help:      "Number of connected peers",
@@ -42,6 +50,7 @@ var (
 
 func init() {
 	prometheus.MustRegister(
+		estimatedNetworkSize,
 		peersConnected,
 		servAndNodeVersion,
 		poolCount,
@@ -49,6 +58,10 @@ func init() {
 	)
 }
 
+func updateNetworkSizeMetric(sz int) {
+	estimatedNetworkSize.Set(float64(sz))
+}
+
 func updateBlockQueueLenMetric(bqLen int) {
 	blockQueueLength.Set(float64(bqLen))
 }