feat: memory based connection limits

Implements memory-based connection limits where the user can specify an absolute or a relative limit of the process' memory usage in relation to the available system memory.

Related: #4252.

Pull-Request: #4281.

Cargo.lock

@@ -2537,6 +2537,7 @@ dependencies = [
  "libp2p-identity",
  "libp2p-kad",
  "libp2p-mdns",
+ "libp2p-memory-connection-limits",
  "libp2p-metrics",
  "libp2p-noise",
  "libp2p-ping",
@@ -2870,6 +2871,24 @@ dependencies = [
  "void",
 ]
 
+[[package]]
+name = "libp2p-memory-connection-limits"
+version = "0.1.0"
+dependencies = [
+ "async-std",
+ "libp2p-core",
+ "libp2p-identify",
+ "libp2p-identity",
+ "libp2p-swarm",
+ "libp2p-swarm-derive",
+ "libp2p-swarm-test",
+ "log",
+ "memory-stats",
+ "rand 0.8.5",
+ "sysinfo",
+ "void",
+]
+
 [[package]]
 name = "libp2p-metrics"
 version = "0.13.1"
@@ -3528,6 +3547,16 @@ dependencies = [
  "autocfg",
 ]
 
+[[package]]
+name = "memory-stats"
+version = "1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "34f79cf9964c5c9545493acda1263f1912f8d2c56c8a2ffee2606cb960acaacc"
+dependencies = [
+ "libc",
+ "winapi",
+]
+
 [[package]]
 name = "metrics-example"
 version = "0.1.0"
@@ -3758,6 +3787,15 @@ dependencies = [
  "minimal-lexical",
 ]
 
+[[package]]
+name = "ntapi"
+version = "0.4.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e8a3895c6391c39d7fe7ebc444a87eb2991b2a0bc718fdabd071eec617fc68e4"
+dependencies = [
+ "winapi",
+]
+
 [[package]]
 name = "nu-ansi-term"
 version = "0.46.0"
@@ -5450,6 +5488,21 @@ dependencies = [
  "unicode-xid",
 ]
 
+[[package]]
+name = "sysinfo"
+version = "0.29.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "165d6d8539689e3d3bc8b98ac59541e1f21c7de7c85d60dc80e43ae0ed2113db"
+dependencies = [
+ "cfg-if 1.0.0",
+ "core-foundation-sys",
+ "libc",
+ "ntapi",
+ "once_cell",
+ "rayon",
+ "winapi",
+]
+
 [[package]]
 name = "system-configuration"
 version = "0.5.1"

Cargo.toml

@@ -18,6 +18,7 @@ members = [
     "misc/allow-block-list",
     "misc/connection-limits",
     "misc/keygen",
+    "misc/memory-connection-limits",
     "misc/metrics",
     "misc/multistream-select",
     "misc/quick-protobuf-codec",
@@ -75,6 +76,7 @@ libp2p-identify = { version = "0.43.0", path = "protocols/identify" }
 libp2p-identity = { version = "0.2.2" }
 libp2p-kad = { version = "0.44.4", path = "protocols/kad" }
 libp2p-mdns = { version = "0.44.0", path = "protocols/mdns" }
+libp2p-memory-connection-limits = { version = "0.1.0", path = "misc/memory-connection-limits" }
 libp2p-metrics = { version = "0.13.1", path = "misc/metrics" }
 libp2p-mplex = { version = "0.40.0", path = "muxers/mplex" }
 libp2p-muxer-test-harness = { path = "muxers/test-harness" }

libp2p/CHANGELOG.md

@@ -6,8 +6,12 @@
 - Add `json` feature which exposes `request_response::json`.
   See [PR 4188].
 
+- Add `libp2p-memory-connection-limits` providing memory usage based connection limit configurations.
+  See [PR 4281].
+
 [PR 4188]: https://github.com/libp2p/rust-libp2p/pull/4188
 [PR 4217]: https://github.com/libp2p/rust-libp2p/pull/4217
+[PR 4281]: https://github.com/libp2p/rust-libp2p/pull/4281
 
 ## 0.52.1
 

libp2p/Cargo.toml

@@ -27,6 +27,7 @@ full = [
     "kad",
     "macros",
     "mdns",
+    "memory-connection-limits",
     "metrics",
     "noise",
     "ping",
@@ -65,6 +66,7 @@ json = ["libp2p-request-response?/json"]
 kad = ["dep:libp2p-kad", "libp2p-metrics?/kad"]
 macros = ["libp2p-swarm/macros"]
 mdns = ["dep:libp2p-mdns"]
+memory-connection-limits = ["dep:libp2p-memory-connection-limits"]
 metrics = ["dep:libp2p-metrics"]
 noise = ["dep:libp2p-noise"]
 ping = ["dep:libp2p-ping", "libp2p-metrics?/ping"]
@@ -124,6 +126,7 @@ pin-project = "1.0.0"
 libp2p-deflate = { workspace = true, optional = true }
 libp2p-dns = { workspace = true, optional = true }
 libp2p-mdns = { workspace = true, optional = true }
+libp2p-memory-connection-limits = { workspace = true, optional = true }
 libp2p-tcp = { workspace = true, optional = true }
 libp2p-tls = { workspace = true, optional = true }
 libp2p-uds = { workspace = true, optional = true }

libp2p/src/lib.rs

@@ -78,6 +78,11 @@ pub use libp2p_kad as kad;
 #[cfg_attr(docsrs, doc(cfg(feature = "mdns")))]
 #[doc(inline)]
 pub use libp2p_mdns as mdns;
+#[cfg(feature = "memory-connection-limits")]
+#[cfg(not(target_arch = "wasm32"))]
+#[cfg_attr(docsrs, doc(cfg(feature = "memory-connection-limits")))]
+#[doc(inline)]
+pub use libp2p_memory_connection_limits as memory_connection_limits;
 #[cfg(feature = "metrics")]
 #[doc(inline)]
 pub use libp2p_metrics as metrics;

misc/memory-connection-limits/CHANGELOG.md

@@ -0,0 +1,3 @@
+## 0.1.0 - unreleased
+
+- Initial release.

misc/memory-connection-limits/Cargo.toml

@@ -0,0 +1,26 @@
+[package]
+name = "libp2p-memory-connection-limits"
+edition = "2021"
+rust-version = { workspace = true }
+description = "Memory usage based connection limits for libp2p."
+version = "0.1.0"
+license = "MIT"
+repository = "https://github.com/libp2p/rust-libp2p"
+keywords = ["peer-to-peer", "libp2p", "networking"]
+categories = ["network-programming", "asynchronous"]
+
+[dependencies]
+memory-stats = { version = "1", features = ["always_use_statm"] }
+libp2p-core = { workspace = true }
+libp2p-swarm = { workspace = true }
+libp2p-identity = { workspace = true, features = ["peerid"] }
+log = "0.4"
+sysinfo = "0.29"
+void = "1"
+
+[dev-dependencies]
+async-std = { version = "1.12.0", features = ["attributes"] }
+libp2p-identify = { workspace = true }
+libp2p-swarm-derive = { path = "../../swarm-derive" }
+libp2p-swarm-test = { path = "../../swarm-test" }
+rand = "0.8.5"

misc/memory-connection-limits/src/lib.rs

@@ -0,0 +1,232 @@
+// Copyright 2023 Protocol Labs.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a
+// copy of this software and associated documentation files (the "Software"),
+// to deal in the Software without restriction, including without limitation
+// the rights to use, copy, modify, merge, publish, distribute, sublicense,
+// and/or sell copies of the Software, and to permit persons to whom the
+// Software is furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
+// OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+// DEALINGS IN THE SOFTWARE.
+
+use libp2p_core::{Endpoint, Multiaddr};
+use libp2p_identity::PeerId;
+use libp2p_swarm::{
+    dummy, ConnectionDenied, ConnectionId, FromSwarm, NetworkBehaviour, PollParameters, THandler,
+    THandlerInEvent, THandlerOutEvent, ToSwarm,
+};
+use void::Void;
+
+use std::{
+    fmt,
+    task::{Context, Poll},
+    time::{Duration, Instant},
+};
+
+/// A [`NetworkBehaviour`] that enforces a set of memory usage based limits.
+///
+/// For these limits to take effect, this needs to be composed into the behaviour tree of your application.
+///
+/// If a connection is denied due to a limit, either a [`SwarmEvent::IncomingConnectionError`](libp2p_swarm::SwarmEvent::IncomingConnectionError)
+/// or [`SwarmEvent::OutgoingConnectionError`](libp2p_swarm::SwarmEvent::OutgoingConnectionError) will be emitted.
+/// The [`ListenError::Denied`](libp2p_swarm::ListenError::Denied) and respectively the [`DialError::Denied`](libp2p_swarm::DialError::Denied) variant
+/// contain a [`ConnectionDenied`](libp2p_swarm::ConnectionDenied) type that can be downcast to [`MemoryUsageLimitExceeded`] error if (and only if) **this**
+/// behaviour denied the connection.
+///
+/// If you employ multiple [`NetworkBehaviour`]s that manage connections, it may also be a different error.
+///
+/// [Behaviour::with_max_bytes] and [Behaviour::with_max_percentage] are mutually exclusive.
+/// If you need to employ both of them, compose two instances of [Behaviour] into your custom behaviour.
+///
+/// # Example
+///
+/// ```rust
+/// # use libp2p_identify as identify;
+/// # use libp2p_swarm_derive::NetworkBehaviour;
+/// # use libp2p_memory_connection_limits as memory_connection_limits;
+///
+/// #[derive(NetworkBehaviour)]
+/// # #[behaviour(prelude = "libp2p_swarm::derive_prelude")]
+/// struct MyBehaviour {
+///   identify: identify::Behaviour,
+///   limits: memory_connection_limits::Behaviour
+/// }
+/// ```
+pub struct Behaviour {
+    max_allowed_bytes: usize,
+    process_physical_memory_bytes: usize,
+    last_refreshed: Instant,
+}
+
+/// The maximum duration for which the retrieved memory-stats of the process are allowed to be stale.
+///
+/// Once exceeded, we will retrieve new stats.
+const MAX_STALE_DURATION: Duration = Duration::from_millis(100);
+
+impl Behaviour {
+    /// Sets the process memory usage threshold in absolute bytes.
+    ///
+    /// New inbound and outbound connections will be denied when the threshold is reached.
+    pub fn with_max_bytes(max_allowed_bytes: usize) -> Self {
+        Self {
+            max_allowed_bytes,
+            process_physical_memory_bytes: memory_stats::memory_stats()
+                .map(|s| s.physical_mem)
+                .unwrap_or_default(),
+            last_refreshed: Instant::now(),
+        }
+    }
+
+    /// Sets the process memory usage threshold in the percentage of the total physical memory.
+    ///
+    /// New inbound and outbound connections will be denied when the threshold is reached.
+    pub fn with_max_percentage(percentage: f64) -> Self {
+        use sysinfo::{RefreshKind, SystemExt};
+
+        let system_memory_bytes =
+            sysinfo::System::new_with_specifics(RefreshKind::new().with_memory()).total_memory();
+
+        Self::with_max_bytes((system_memory_bytes as f64 * percentage).round() as usize)
+    }
+
+    /// Gets the process memory usage threshold in bytes.
+    pub fn max_allowed_bytes(&self) -> usize {
+        self.max_allowed_bytes
+    }
+
+    fn check_limit(&mut self) -> Result<(), ConnectionDenied> {
+        self.refresh_memory_stats_if_needed();
+
+        if self.process_physical_memory_bytes > self.max_allowed_bytes {
+            return Err(ConnectionDenied::new(MemoryUsageLimitExceeded {
+                process_physical_memory_bytes: self.process_physical_memory_bytes,
+                max_allowed_bytes: self.max_allowed_bytes,
+            }));
+        }
+
+        Ok(())
+    }
+
+    fn refresh_memory_stats_if_needed(&mut self) {
+        let now = Instant::now();
+
+        if self.last_refreshed + MAX_STALE_DURATION > now {
+            // Memory stats are reasonably recent, don't refresh.
+            return;
+        }
+
+        let stats = match memory_stats::memory_stats() {
+            Some(stats) => stats,
+            None => {
+                log::warn!("Failed to retrieve process memory stats");
+                return;
+            }
+        };
+
+        self.last_refreshed = now;
+        self.process_physical_memory_bytes = stats.physical_mem;
+    }
+}
+
+impl NetworkBehaviour for Behaviour {
+    type ConnectionHandler = dummy::ConnectionHandler;
+    type ToSwarm = Void;
+
+    fn handle_pending_inbound_connection(
+        &mut self,
+        _: ConnectionId,
+        _: &Multiaddr,
+        _: &Multiaddr,
+    ) -> Result<(), ConnectionDenied> {
+        self.check_limit()
+    }
+
+    fn handle_established_inbound_connection(
+        &mut self,
+        _: ConnectionId,
+        _: PeerId,
+        _: &Multiaddr,
+        _: &Multiaddr,
+    ) -> Result<THandler<Self>, ConnectionDenied> {
+        Ok(dummy::ConnectionHandler)
+    }
+
+    fn handle_pending_outbound_connection(
+        &mut self,
+        _: ConnectionId,
+        _: Option<PeerId>,
+        _: &[Multiaddr],
+        _: Endpoint,
+    ) -> Result<Vec<Multiaddr>, ConnectionDenied> {
+        self.check_limit()?;
+        Ok(vec![])
+    }
+
+    fn handle_established_outbound_connection(
+        &mut self,
+        _: ConnectionId,
+        _: PeerId,
+        _: &Multiaddr,
+        _: Endpoint,
+    ) -> Result<THandler<Self>, ConnectionDenied> {
+        Ok(dummy::ConnectionHandler)
+    }
+
+    fn on_swarm_event(&mut self, _: FromSwarm<Self::ConnectionHandler>) {}
+
+    fn on_connection_handler_event(
+        &mut self,
+        _id: PeerId,
+        _: ConnectionId,
+        event: THandlerOutEvent<Self>,
+    ) {
+        void::unreachable(event)
+    }
+
+    fn poll(
+        &mut self,
+        _: &mut Context<'_>,
+        _: &mut impl PollParameters,
+    ) -> Poll<ToSwarm<Self::ToSwarm, THandlerInEvent<Self>>> {
+        Poll::Pending
+    }
+}
+
+/// A connection limit has been exceeded.
+#[derive(Debug, Clone, Copy)]
+pub struct MemoryUsageLimitExceeded {
+    process_physical_memory_bytes: usize,
+    max_allowed_bytes: usize,
+}
+
+impl MemoryUsageLimitExceeded {
+    pub fn process_physical_memory_bytes(&self) -> usize {
+        self.process_physical_memory_bytes
+    }
+
+    pub fn max_allowed_bytes(&self) -> usize {
+        self.max_allowed_bytes
+    }
+}
+
+impl std::error::Error for MemoryUsageLimitExceeded {}
+
+impl fmt::Display for MemoryUsageLimitExceeded {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(
+            f,
+            "process physical memory usage limit exceeded: process memory: {} bytes, max allowed: {} bytes",
+            self.process_physical_memory_bytes,
+            self.max_allowed_bytes,
+        )
+    }
+}

misc/memory-connection-limits/tests/max_bytes.rs

@@ -0,0 +1,89 @@
+// Copyright 2023 Protocol Labs.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a
+// copy of this software and associated documentation files (the "Software"),
+// to deal in the Software without restriction, including without limitation
+// the rights to use, copy, modify, merge, publish, distribute, sublicense,
+// and/or sell copies of the Software, and to permit persons to whom the
+// Software is furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
+// OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+// DEALINGS IN THE SOFTWARE.
+
+mod util;
+
+use libp2p_core::Multiaddr;
+use libp2p_identity::PeerId;
+use libp2p_memory_connection_limits::*;
+use std::time::Duration;
+use util::*;
+
+use libp2p_swarm::{dial_opts::DialOpts, DialError, Swarm};
+use libp2p_swarm_test::SwarmExt;
+
+#[test]
+fn max_bytes() {
+    const CONNECTION_LIMIT: usize = 20;
+    let max_allowed_bytes = CONNECTION_LIMIT * 1024 * 1024;
+
+    let mut network = Swarm::new_ephemeral(|_| TestBehaviour {
+        connection_limits: Behaviour::with_max_bytes(max_allowed_bytes),
+        mem_consumer: ConsumeMemoryBehaviour1MBPending0Established::default(),
+    });
+
+    let addr: Multiaddr = "/memory/1234".parse().unwrap();
+    let target = PeerId::random();
+
+    // Exercise `dial` function to get more stable memory stats later
+    network
+        .dial(
+            DialOpts::peer_id(target)
+                .addresses(vec![addr.clone()])
+                .build(),
+        )
+        .expect("Unexpected connection limit.");
+
+    // Adds current mem usage to the limit and update
+    let max_allowed_bytes_plus_base_usage =
+        max_allowed_bytes + memory_stats::memory_stats().unwrap().physical_mem;
+    network.behaviour_mut().connection_limits =
+        Behaviour::with_max_bytes(max_allowed_bytes_plus_base_usage);
+
+    for _ in 0..CONNECTION_LIMIT {
+        network
+            .dial(
+                DialOpts::peer_id(target)
+                    .addresses(vec![addr.clone()])
+                    .build(),
+            )
+            .expect("Unexpected connection limit.");
+    }
+
+    std::thread::sleep(Duration::from_millis(100)); // Memory stats are only updated every 100ms internally, ensure they are up-to-date when we try to exceed it.
+
+    match network
+        .dial(DialOpts::peer_id(target).addresses(vec![addr]).build())
+        .expect_err("Unexpected dialing success.")
+    {
+        DialError::Denied { cause } => {
+            let exceeded = cause
+                .downcast::<MemoryUsageLimitExceeded>()
+                .expect("connection denied because of limit");
+
+            assert_eq!(
+                exceeded.max_allowed_bytes(),
+                max_allowed_bytes_plus_base_usage
+            );
+            assert!(exceeded.process_physical_memory_bytes() >= exceeded.max_allowed_bytes());
+        }
+        e => panic!("Unexpected error: {e:?}"),
+    }
+}

misc/memory-connection-limits/tests/max_percentage.rs

@@ -0,0 +1,88 @@
+// Copyright 2023 Protocol Labs.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a
+// copy of this software and associated documentation files (the "Software"),
+// to deal in the Software without restriction, including without limitation
+// the rights to use, copy, modify, merge, publish, distribute, sublicense,
+// and/or sell copies of the Software, and to permit persons to whom the
+// Software is furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
+// OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+// DEALINGS IN THE SOFTWARE.
+
+mod util;
+
+use libp2p_core::Multiaddr;
+use libp2p_identity::PeerId;
+use libp2p_memory_connection_limits::*;
+use std::time::Duration;
+use sysinfo::{RefreshKind, SystemExt};
+use util::*;
+
+use libp2p_swarm::{dial_opts::DialOpts, DialError, Swarm};
+use libp2p_swarm_test::SwarmExt;
+
+#[test]
+fn max_percentage() {
+    const CONNECTION_LIMIT: usize = 20;
+    let system_info = sysinfo::System::new_with_specifics(RefreshKind::new().with_memory());
+
+    let mut network = Swarm::new_ephemeral(|_| TestBehaviour {
+        connection_limits: Behaviour::with_max_percentage(0.1),
+        mem_consumer: ConsumeMemoryBehaviour1MBPending0Established::default(),
+    });
+
+    let addr: Multiaddr = "/memory/1234".parse().unwrap();
+    let target = PeerId::random();
+
+    // Exercise `dial` function to get more stable memory stats later
+    network
+        .dial(
+            DialOpts::peer_id(target)
+                .addresses(vec![addr.clone()])
+                .build(),
+        )
+        .expect("Unexpected connection limit.");
+
+    // Adds current mem usage to the limit and update
+    let current_mem = memory_stats::memory_stats().unwrap().physical_mem;
+    let max_allowed_bytes = current_mem + CONNECTION_LIMIT * 1024 * 1024;
+    network.behaviour_mut().connection_limits = Behaviour::with_max_percentage(
+        max_allowed_bytes as f64 / system_info.total_memory() as f64,
+    );
+
+    for _ in 0..CONNECTION_LIMIT {
+        network
+            .dial(
+                DialOpts::peer_id(target)
+                    .addresses(vec![addr.clone()])
+                    .build(),
+            )
+            .expect("Unexpected connection limit.");
+    }
+
+    std::thread::sleep(Duration::from_millis(100)); // Memory stats are only updated every 100ms internally, ensure they are up-to-date when we try to exceed it.
+
+    match network
+        .dial(DialOpts::peer_id(target).addresses(vec![addr]).build())
+        .expect_err("Unexpected dialing success.")
+    {
+        DialError::Denied { cause } => {
+            let exceeded = cause
+                .downcast::<MemoryUsageLimitExceeded>()
+                .expect("connection denied because of limit");
+
+            assert_eq!(exceeded.max_allowed_bytes(), max_allowed_bytes);
+            assert!(exceeded.process_physical_memory_bytes() >= exceeded.max_allowed_bytes());
+        }
+        e => panic!("Unexpected error: {e:?}"),
+    }
+}

misc/memory-connection-limits/tests/util.rs

@@ -0,0 +1,128 @@
+// Copyright 2023 Protocol Labs.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a
+// copy of this software and associated documentation files (the "Software"),
+// to deal in the Software without restriction, including without limitation
+// the rights to use, copy, modify, merge, publish, distribute, sublicense,
+// and/or sell copies of the Software, and to permit persons to whom the
+// Software is furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
+// OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+// DEALINGS IN THE SOFTWARE.
+
+use std::task::{Context, Poll};
+
+use libp2p_core::{Endpoint, Multiaddr};
+use libp2p_identity::PeerId;
+use libp2p_swarm::{
+    dummy, ConnectionDenied, ConnectionId, FromSwarm, NetworkBehaviour, PollParameters, THandler,
+    THandlerInEvent, THandlerOutEvent, ToSwarm,
+};
+use void::Void;
+
+#[derive(libp2p_swarm_derive::NetworkBehaviour)]
+#[behaviour(prelude = "libp2p_swarm::derive_prelude")]
+pub(crate) struct TestBehaviour {
+    pub(crate) connection_limits: libp2p_memory_connection_limits::Behaviour,
+    pub(crate) mem_consumer: ConsumeMemoryBehaviour1MBPending0Established,
+}
+
+pub(crate) type ConsumeMemoryBehaviour1MBPending0Established =
+    ConsumeMemoryBehaviour<{ 1024 * 1024 }, 0>;
+
+#[derive(Default)]
+pub(crate) struct ConsumeMemoryBehaviour<const MEM_PENDING: usize, const MEM_ESTABLISHED: usize> {
+    mem_pending: Vec<Vec<u8>>,
+    mem_established: Vec<Vec<u8>>,
+}
+
+impl<const MEM_PENDING: usize, const MEM_ESTABLISHED: usize>
+    ConsumeMemoryBehaviour<MEM_PENDING, MEM_ESTABLISHED>
+{
+    fn handle_pending(&mut self) {
+        // 1MB
+        self.mem_pending.push(vec![1; MEM_PENDING]);
+    }
+
+    fn handle_established(&mut self) {
+        // 1MB
+        self.mem_established.push(vec![1; MEM_ESTABLISHED]);
+    }
+}
+
+impl<const MEM_PENDING: usize, const MEM_ESTABLISHED: usize> NetworkBehaviour
+    for ConsumeMemoryBehaviour<MEM_PENDING, MEM_ESTABLISHED>
+{
+    type ConnectionHandler = dummy::ConnectionHandler;
+    type ToSwarm = Void;
+
+    fn handle_pending_inbound_connection(
+        &mut self,
+        _: ConnectionId,
+        _: &Multiaddr,
+        _: &Multiaddr,
+    ) -> Result<(), ConnectionDenied> {
+        self.handle_pending();
+        Ok(())
+    }
+
+    fn handle_pending_outbound_connection(
+        &mut self,
+        _: ConnectionId,
+        _: Option<PeerId>,
+        _: &[Multiaddr],
+        _: Endpoint,
+    ) -> Result<Vec<Multiaddr>, ConnectionDenied> {
+        self.handle_pending();
+        Ok(vec![])
+    }
+
+    fn handle_established_inbound_connection(
+        &mut self,
+        _: ConnectionId,
+        _: PeerId,
+        _: &Multiaddr,
+        _: &Multiaddr,
+    ) -> Result<THandler<Self>, ConnectionDenied> {
+        self.handle_established();
+        Ok(dummy::ConnectionHandler)
+    }
+
+    fn handle_established_outbound_connection(
+        &mut self,
+        _: ConnectionId,
+        _: PeerId,
+        _: &Multiaddr,
+        _: Endpoint,
+    ) -> Result<THandler<Self>, ConnectionDenied> {
+        self.handle_established();
+        Ok(dummy::ConnectionHandler)
+    }
+
+    fn on_swarm_event(&mut self, _: FromSwarm<Self::ConnectionHandler>) {}
+
+    fn on_connection_handler_event(
+        &mut self,
+        _id: PeerId,
+        _: ConnectionId,
+        event: THandlerOutEvent<Self>,
+    ) {
+        void::unreachable(event)
+    }
+
+    fn poll(
+        &mut self,
+        _: &mut Context<'_>,
+        _: &mut impl PollParameters,
+    ) -> Poll<ToSwarm<Self::ToSwarm, THandlerInEvent<Self>>> {
+        Poll::Pending
+    }
+}