p2p/discover: simplify table meterics tracking

Author: lightclient

p2p/discover/table.go

@@ -81,7 +81,8 @@ type Table struct {
 	closeReq   chan struct{}
 	closed     chan struct{}
 
-	nodeAddedHook func(*node) // for testing
+	nodeAddedHook   func(*bucket, *node)
+	nodeRemovedHook func(*bucket, *node)
 }
 
 // transport is implemented by the UDP transports.
@@ -131,6 +132,22 @@ func newTable(t transport, db *enode.DB, cfg Config) (*Table, error) {
 	return tab, nil
 }
 
+func newMeteredTable(t transport, db *enode.DB, cfg Config) (*Table, error) {
+	tab, err := newTable(t, db, cfg)
+	if err != nil {
+		return nil, err
+	}
+	if metrics.Enabled {
+		tab.nodeAddedHook = func(b *bucket, n *node) {
+			bucketsCounter[b.index].Inc(1)
+		}
+		tab.nodeRemovedHook = func(b *bucket, n *node) {
+			bucketsCounter[b.index].Dec(1)
+		}
+	}
+	return tab, nil
+}
+
 // Nodes returns all nodes contained in the table.
 func (tab *Table) Nodes() []*enode.Node {
 	if !tab.isInitDone() {
@@ -498,10 +515,7 @@ func (tab *Table) addSeenNode(n *node) {
 	n.addedAt = time.Now()
 
 	if tab.nodeAddedHook != nil {
-		tab.nodeAddedHook(n)
-	}
-	if metrics.Enabled {
-		bucketsCounter[b.index].Inc(1)
+		tab.nodeAddedHook(b, n)
 	}
 }
 
@@ -545,10 +559,7 @@ func (tab *Table) addVerifiedNode(n *node) {
 	n.addedAt = time.Now()
 
 	if tab.nodeAddedHook != nil {
-		tab.nodeAddedHook(n)
-	}
-	if metrics.Enabled {
-		bucketsCounter[b.index].Inc(1)
+		tab.nodeAddedHook(b, n)
 	}
 }
 
@@ -647,11 +658,16 @@ func (tab *Table) bumpInBucket(b *bucket, n *node) bool {
 }
 
 func (tab *Table) deleteInBucket(b *bucket, n *node) {
-	if metrics.Enabled && contains(b.entries, n.ID()) {
-		bucketsCounter[b.index].Dec(1)
+	// Check if the node is actually in the bucket so the removed hook
+	// isn't called multiple times for the same node.
+	if !contains(b.entries, n.ID()) {
+		return
 	}
 	b.entries = deleteNode(b.entries, n)
 	tab.removeIP(b, n.IP())
+	if tab.nodeRemovedHook != nil {
+		tab.nodeRemovedHook(b, n)
+	}
 }
 
 func contains(ns []*node, id enode.ID) bool {

p2p/discover/v4_udp.go

@@ -142,7 +142,7 @@ func ListenV4(c UDPConn, ln *enode.LocalNode, cfg Config) (*UDPv4, error) {
 		log:             cfg.Log,
 	}
 
-	tab, err := newTable(t, ln.Database(), cfg)
+	tab, err := newMeteredTable(t, ln.Database(), cfg)
 	if err != nil {
 		return nil, err
 	}

p2p/discover/v4_udp_test.go

@@ -394,7 +394,7 @@ func TestUDPv4_pingMatchIP(t *testing.T) {
 func TestUDPv4_successfulPing(t *testing.T) {
 	test := newUDPTest(t)
 	added := make(chan *node, 1)
-	test.table.nodeAddedHook = func(n *node) { added <- n }
+	test.table.nodeAddedHook = func(b *bucket, n *node) { added <- n }
 	defer test.close()
 
 	// The remote side sends a ping packet to initiate the exchange.

p2p/discover/v5_udp.go

@@ -174,7 +174,7 @@ func newUDPv5(conn UDPConn, ln *enode.LocalNode, cfg Config) (*UDPv5, error) {
 		cancelCloseCtx: cancelCloseCtx,
 	}
 	t.talk = newTalkSystem(t)
-	tab, err := newTable(t, t.db, cfg)
+	tab, err := newMeteredTable(t, t.db, cfg)
 	if err != nil {
 		return nil, err
 	}