From b5d993267f7d76353ede4ddb2879bfe757851b1e Mon Sep 17 00:00:00 2001 From: hanabi1224 Date: Wed, 9 Aug 2023 00:29:21 +0800 Subject: [PATCH] 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 | 53 ++++ Cargo.toml | 2 + libp2p/CHANGELOG.md | 4 + libp2p/Cargo.toml | 3 + libp2p/src/lib.rs | 5 + misc/memory-connection-limits/CHANGELOG.md | 3 + misc/memory-connection-limits/Cargo.toml | 26 ++ misc/memory-connection-limits/src/lib.rs | 232 ++++++++++++++++++ .../tests/max_bytes.rs | 89 +++++++ .../tests/max_percentage.rs | 88 +++++++ misc/memory-connection-limits/tests/util.rs | 128 ++++++++++ 11 files changed, 633 insertions(+) create mode 100644 misc/memory-connection-limits/CHANGELOG.md create mode 100644 misc/memory-connection-limits/Cargo.toml create mode 100644 misc/memory-connection-limits/src/lib.rs create mode 100644 misc/memory-connection-limits/tests/max_bytes.rs create mode 100644 misc/memory-connection-limits/tests/max_percentage.rs create mode 100644 misc/memory-connection-limits/tests/util.rs diff --git a/Cargo.lock b/Cargo.lock index 547efe4391a..2cf8332a7cd 100644 --- a/Cargo.lock +++ b/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" diff --git a/Cargo.toml b/Cargo.toml index 702a847b0c2..b2d30c9257d 100644 --- a/Cargo.toml +++ b/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" } diff --git a/libp2p/CHANGELOG.md b/libp2p/CHANGELOG.md index 4161ed054c5..fc2b24bd6e6 100644 --- a/libp2p/CHANGELOG.md +++ b/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 diff --git a/libp2p/Cargo.toml b/libp2p/Cargo.toml index fd9290cf506..e057c9e8dd4 100644 --- a/libp2p/Cargo.toml +++ b/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 } diff --git a/libp2p/src/lib.rs b/libp2p/src/lib.rs index d22e4c25d30..9ec1fe48af4 100644 --- a/libp2p/src/lib.rs +++ b/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; diff --git a/misc/memory-connection-limits/CHANGELOG.md b/misc/memory-connection-limits/CHANGELOG.md new file mode 100644 index 00000000000..5dac7afbf35 --- /dev/null +++ b/misc/memory-connection-limits/CHANGELOG.md @@ -0,0 +1,3 @@ +## 0.1.0 - unreleased + +- Initial release. diff --git a/misc/memory-connection-limits/Cargo.toml b/misc/memory-connection-limits/Cargo.toml new file mode 100644 index 00000000000..6bfecfd2b0e --- /dev/null +++ b/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" diff --git a/misc/memory-connection-limits/src/lib.rs b/misc/memory-connection-limits/src/lib.rs new file mode 100644 index 00000000000..33e40b11843 --- /dev/null +++ b/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, ConnectionDenied> { + Ok(dummy::ConnectionHandler) + } + + fn handle_pending_outbound_connection( + &mut self, + _: ConnectionId, + _: Option, + _: &[Multiaddr], + _: Endpoint, + ) -> Result, ConnectionDenied> { + self.check_limit()?; + Ok(vec![]) + } + + fn handle_established_outbound_connection( + &mut self, + _: ConnectionId, + _: PeerId, + _: &Multiaddr, + _: Endpoint, + ) -> Result, ConnectionDenied> { + Ok(dummy::ConnectionHandler) + } + + fn on_swarm_event(&mut self, _: FromSwarm) {} + + fn on_connection_handler_event( + &mut self, + _id: PeerId, + _: ConnectionId, + event: THandlerOutEvent, + ) { + void::unreachable(event) + } + + fn poll( + &mut self, + _: &mut Context<'_>, + _: &mut impl PollParameters, + ) -> Poll>> { + 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, + ) + } +} diff --git a/misc/memory-connection-limits/tests/max_bytes.rs b/misc/memory-connection-limits/tests/max_bytes.rs new file mode 100644 index 00000000000..af86b048785 --- /dev/null +++ b/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::() + .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:?}"), + } +} diff --git a/misc/memory-connection-limits/tests/max_percentage.rs b/misc/memory-connection-limits/tests/max_percentage.rs new file mode 100644 index 00000000000..ea3f20e6cbc --- /dev/null +++ b/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::() + .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:?}"), + } +} diff --git a/misc/memory-connection-limits/tests/util.rs b/misc/memory-connection-limits/tests/util.rs new file mode 100644 index 00000000000..a2fd7c20fed --- /dev/null +++ b/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 { + mem_pending: Vec>, + mem_established: Vec>, +} + +impl + ConsumeMemoryBehaviour +{ + 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 NetworkBehaviour + for ConsumeMemoryBehaviour +{ + 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, + _: &[Multiaddr], + _: Endpoint, + ) -> Result, ConnectionDenied> { + self.handle_pending(); + Ok(vec![]) + } + + fn handle_established_inbound_connection( + &mut self, + _: ConnectionId, + _: PeerId, + _: &Multiaddr, + _: &Multiaddr, + ) -> Result, ConnectionDenied> { + self.handle_established(); + Ok(dummy::ConnectionHandler) + } + + fn handle_established_outbound_connection( + &mut self, + _: ConnectionId, + _: PeerId, + _: &Multiaddr, + _: Endpoint, + ) -> Result, ConnectionDenied> { + self.handle_established(); + Ok(dummy::ConnectionHandler) + } + + fn on_swarm_event(&mut self, _: FromSwarm) {} + + fn on_connection_handler_event( + &mut self, + _id: PeerId, + _: ConnectionId, + event: THandlerOutEvent, + ) { + void::unreachable(event) + } + + fn poll( + &mut self, + _: &mut Context<'_>, + _: &mut impl PollParameters, + ) -> Poll>> { + Poll::Pending + } +}