From 2e111bd968385d4ecc07acd595dcf08ba149d19c Mon Sep 17 00:00:00 2001
From: Alexis Sellier <self@cloudhead.io>
Date: Sat, 11 Nov 2023 13:26:38 +0100
Subject: [PATCH] [p2p] Use event for peer disconnection

---
 p2p/src/fsm.rs         |  5 -----
 p2p/src/fsm/addrmgr.rs |  5 ++++-
 p2p/src/fsm/cbfmgr.rs  |  8 +++-----
 p2p/src/fsm/invmgr.rs  | 10 ++++------
 p2p/src/fsm/pingmgr.rs | 10 ++++------
 p2p/src/fsm/syncmgr.rs |  8 +++-----
 6 files changed, 18 insertions(+), 28 deletions(-)

diff --git a/p2p/src/fsm.rs b/p2p/src/fsm.rs
index 014f1c4d..adf8b20c 100644
--- a/p2p/src/fsm.rs
+++ b/p2p/src/fsm.rs
@@ -813,13 +813,8 @@ impl<T: BlockTree, F: Filters, P: peer::Store, C: AdjustedClock<PeerId>> traits:
         addr: &net::SocketAddr,
         reason: nakamoto_net::Disconnect<DisconnectReason>,
     ) {
-        self.cbfmgr.peer_disconnected(addr);
-        self.syncmgr.peer_disconnected(addr);
-        self.addrmgr.peer_disconnected(addr, reason.clone());
-        self.pingmgr.peer_disconnected(addr);
         self.peermgr
             .peer_disconnected(addr, &mut self.addrmgr, reason);
-        self.invmgr.peer_disconnected(addr);
     }
 
     fn tick(&mut self, local_time: LocalTime) {
diff --git a/p2p/src/fsm/addrmgr.rs b/p2p/src/fsm/addrmgr.rs
index e6c13f71..7ada8c75 100644
--- a/p2p/src/fsm/addrmgr.rs
+++ b/p2p/src/fsm/addrmgr.rs
@@ -151,6 +151,9 @@ impl<P: Store, C: Clock> AddressManager<P, C> {
                 }
                 self.peer_negotiated(&addr, services, link);
             }
+            Event::PeerDisconnected { addr, reason } => {
+                self.peer_disconnected(&addr, reason);
+            }
             Event::MessageReceived { from, message } => {
                 if let Some(ka) = self.peers.get_mut(&from.ip()) {
                     ka.last_active = Some(self.clock.local_time());
@@ -255,7 +258,7 @@ impl<P: Store, C: Clock> AddressManager<P, C> {
     }
 
     /// Called when a peer disconnected.
-    pub fn peer_disconnected(
+    fn peer_disconnected(
         &mut self,
         addr: &net::SocketAddr,
         reason: Disconnect<super::DisconnectReason>,
diff --git a/p2p/src/fsm/cbfmgr.rs b/p2p/src/fsm/cbfmgr.rs
index b1f0150e..ca02e42b 100644
--- a/p2p/src/fsm/cbfmgr.rs
+++ b/p2p/src/fsm/cbfmgr.rs
@@ -184,6 +184,9 @@ impl<F: Filters, C: Clock> FilterManager<F, C> {
             } => {
                 self.peer_negotiated(addr, height, services, link, persistent, tree);
             }
+            Event::PeerDisconnected { addr, .. } => {
+                self.peers.remove(&addr);
+            }
             Event::BlockProcessed { block, height, .. } => {
                 if self.pending_blocks.remove(&height) {
                     self.outbox.event(Event::BlockMatched { block, height });
@@ -507,11 +510,6 @@ impl<F: Filters, C: Clock> FilterManager<F, C> {
         Ok(())
     }
 
-    /// Called when a peer disconnected.
-    pub fn peer_disconnected(&mut self, id: &PeerId) {
-        self.peers.remove(id);
-    }
-
     /// Called when a new peer was negotiated.
     pub fn peer_negotiated<T: BlockReader>(
         &mut self,
diff --git a/p2p/src/fsm/invmgr.rs b/p2p/src/fsm/invmgr.rs
index 769d6bac..5c88f393 100644
--- a/p2p/src/fsm/invmgr.rs
+++ b/p2p/src/fsm/invmgr.rs
@@ -168,6 +168,9 @@ impl<C: Clock> InventoryManager<C> {
             } => {
                 self.peer_negotiated(addr, services, relay, wtxid_relay);
             }
+            Event::PeerDisconnected { addr, .. } => {
+                self.peers.remove(&addr);
+            }
             Event::BlockHeadersImported {
                 result: ImportResult::TipChanged { reverted, .. },
                 ..
@@ -191,7 +194,7 @@ impl<C: Clock> InventoryManager<C> {
     }
 
     /// Called when a peer is negotiated.
-    pub fn peer_negotiated(
+    fn peer_negotiated(
         &mut self,
         addr: PeerId,
         services: ServiceFlags,
@@ -218,11 +221,6 @@ impl<C: Clock> InventoryManager<C> {
         );
     }
 
-    /// Called when a peer disconnected.
-    pub fn peer_disconnected(&mut self, id: &PeerId) {
-        self.peers.remove(id);
-    }
-
     /// Called when a block is reverted.
     pub fn block_reverted(&mut self, height: Height) {
         self.estimator.rollback(height - 1);
diff --git a/p2p/src/fsm/pingmgr.rs b/p2p/src/fsm/pingmgr.rs
index c6060e44..9c2db489 100644
--- a/p2p/src/fsm/pingmgr.rs
+++ b/p2p/src/fsm/pingmgr.rs
@@ -95,6 +95,9 @@ impl<C: Clock> PingManager<C> {
             Event::PeerNegotiated { addr, .. } => {
                 self.peer_negotiated(addr);
             }
+            Event::PeerDisconnected { addr, .. } => {
+                self.peers.remove(&addr);
+            }
             Event::MessageReceived { from, message } => match message.as_ref() {
                 NetworkMessage::Ping(nonce) => {
                     self.received_ping(from, *nonce);
@@ -109,7 +112,7 @@ impl<C: Clock> PingManager<C> {
     }
 
     /// Called when a peer is negotiated.
-    pub fn peer_negotiated(&mut self, address: PeerId) {
+    fn peer_negotiated(&mut self, address: PeerId) {
         let nonce = self.rng.u64(..);
         let now = self.clock.local_time();
 
@@ -124,11 +127,6 @@ impl<C: Clock> PingManager<C> {
         );
     }
 
-    /// Called when a peer is disconnected.
-    pub fn peer_disconnected(&mut self, addr: &PeerId) {
-        self.peers.remove(addr);
-    }
-
     /// Called when a tick is received.
     pub fn received_wake(&mut self) {
         let now = self.clock.local_time();
diff --git a/p2p/src/fsm/syncmgr.rs b/p2p/src/fsm/syncmgr.rs
index 5c071e7d..d85bee97 100644
--- a/p2p/src/fsm/syncmgr.rs
+++ b/p2p/src/fsm/syncmgr.rs
@@ -162,6 +162,9 @@ impl<C: Clock> SyncManager<C> {
             } => {
                 self.peer_negotiated(addr, height, services, link, tree);
             }
+            Event::PeerDisconnected { addr, .. } => {
+                self.unregister(&addr);
+            }
             Event::MessageReceived { from, message } => match message.as_ref() {
                 NetworkMessage::Headers(headers) => {
                     self.received_headers(&from, headers, tree);
@@ -214,11 +217,6 @@ impl<C: Clock> SyncManager<C> {
         self.sync(tree);
     }
 
-    /// Called when a peer disconnected.
-    pub fn peer_disconnected(&mut self, id: &PeerId) {
-        self.unregister(id);
-    }
-
     /// Called when we received a `getheaders` message from a peer.
     pub fn received_getheaders<T: BlockReader>(
         &mut self,