From 7a5e87e9deba0eeb990618ecc3446a9e7af638d9 Mon Sep 17 00:00:00 2001
From: Divma <26765164+divagant-martian@users.noreply.github.com>
Date: Wed, 9 Feb 2022 10:08:28 -0500
Subject: [PATCH] swarm/src/behaviour: Merge inject_* paired methods (#2445)

Co-authored-by: Max Inden <mail@max-inden.de>
---
 CHANGELOG.md     |  6 +++++-
 src/behaviour.rs | 30 ++++++++++++++++--------------
 2 files changed, 21 insertions(+), 15 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index c93d77b8c06..5df9d722fa0 100644
--- 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.
\ No newline at end of file
+- Initial release.
diff --git a/src/behaviour.rs b/src/behaviour.rs
index a1a8bb6ff21..b2e13b11dde 100644
--- 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,