swarm/src/behaviour: Merge inject_* paired methods (libp2p#2445)

Co-authored-by: Max Inden <mail@max-inden.de>
umgefahren · Feb 9, 2022 · 7a5e87e · 7a5e87e
1 parent 83a1ff6
commit 7a5e87e
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 15 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -6,6 +6,10 @@
 
 - Update to `libp2p-request-response` `v0.16.0`.
 
+- Merge NetworkBehaviour's inject_\* paired methods (see PR 2445).
+
+[PR 2445]: https://github.com/libp2p/rust-libp2p/pull/2445
+
 # 0.1.0 [2022-01-27]
 
-- Initial release.
+- Initial release.
diff --git a/src/behaviour.rs b/src/behaviour.rs
@@ -303,9 +303,15 @@ impl NetworkBehaviour for Behaviour {
         conn: &ConnectionId,
         endpoint: &ConnectedPoint,
         failed_addresses: Option<&Vec<Multiaddr>>,
+        other_established: usize,
     ) {
-        self.inner
-            .inject_connection_established(peer, conn, endpoint, failed_addresses);
+        self.inner.inject_connection_established(
+            peer,
+            conn,
+            endpoint,
+            failed_addresses,
+            other_established,
+        );
         let connections = self.connected.entry(*peer).or_default();
         let addr = if endpoint.is_relayed() {
             None
@@ -342,11 +348,16 @@ impl NetworkBehaviour for Behaviour {
         conn: &ConnectionId,
         endpoint: &ConnectedPoint,
         handler: <Self::ProtocolsHandler as IntoProtocolsHandler>::Handler,
+        remaining_established: usize,
     ) {
         self.inner
-            .inject_connection_closed(peer, conn, endpoint, handler);
-        let connections = self.connected.get_mut(peer).expect("Peer is connected.");
-        connections.remove(conn);
+            .inject_connection_closed(peer, conn, endpoint, handler, remaining_established);
+        if remaining_established == 0 {
+            self.connected.remove(peer);
+        } else {
+            let connections = self.connected.get_mut(peer).expect("Peer is connected.");
+            connections.remove(conn);
+        }
     }
 
     fn inject_dial_failure(
@@ -362,11 +373,6 @@ impl NetworkBehaviour for Behaviour {
         }
     }
 
-    fn inject_disconnected(&mut self, peer: &PeerId) {
-        self.inner.inject_disconnected(peer);
-        self.connected.remove(peer);
-    }
-
     fn inject_address_change(
         &mut self,
         peer: &PeerId,
@@ -461,10 +467,6 @@ impl NetworkBehaviour for Behaviour {
         self.inner.addresses_of_peer(peer)
     }
 
-    fn inject_connected(&mut self, peer: &PeerId) {
-        self.inner.inject_connected(peer)
-    }
-
     fn inject_event(
         &mut self,
         peer_id: PeerId,