feat: memory based connection limits

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" }

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 = "1"
+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/connection_limits.rs

@@ -0,0 +1,96 @@
+// 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 super::*;
+use libp2p_swarm::ConnectionDenied;
+
+/// The configurable memory usage based connection limits.
+#[derive(Debug)]
+pub(crate) struct MemoryUsageBasedConnectionLimits {
+    max_process_memory_usage_bytes: Option<usize>,
+    max_process_memory_usage_percentage: Option<f64>,
+    system_physical_memory_bytes: usize,
+}
+
+impl MemoryUsageBasedConnectionLimits {
+    pub(crate) fn new() -> Self {
+        use sysinfo::{RefreshKind, SystemExt};
+
+        let system_info = sysinfo::System::new_with_specifics(RefreshKind::new().with_memory());
+
+        MemoryUsageBasedConnectionLimits {
+            max_process_memory_usage_bytes: None,
+            max_process_memory_usage_percentage: None,
+            system_physical_memory_bytes: system_info.total_memory() as usize,
+        }
+    }
+
+    /// Sets the process memory usage threshold in bytes,
+    /// all pending connections will be dropped when the threshold is exeeded
+    pub(crate) fn with_max_bytes(mut self, bytes: usize) -> Self {
+        self.max_process_memory_usage_bytes = Some(bytes);
+        self
+    }
+
+    /// Sets the process memory usage threshold in the percentage of the total physical memory,
+    /// all pending connections will be dropped when the threshold is exeeded
+    pub(crate) fn with_max_percentage(mut self, percentage: f64) -> Self {
+        self.max_process_memory_usage_percentage = Some(percentage);
+        self
+    }
+
+    pub(crate) fn check_limit(
+        &self,
+        mem_tracker: &mut ProcessMemoryUsageTracker,
+    ) -> Result<(), ConnectionDenied> {
+        if let Some(max_allowed_bytes) = self.max_allowed_bytes() {
+            mem_tracker.refresh_if_needed();
+            if mem_tracker.physical_bytes > max_allowed_bytes {
+                return Err(ConnectionDenied::new(MemoryUsageLimitExceeded {
+                    process_physical_memory_bytes: mem_tracker.physical_bytes,
+                    max_allowed_bytes,
+                }));
+            }
+        }
+
+        Ok(())
+    }
+
+    fn max_allowed_bytes(&self) -> Option<usize> {
+        let max_process_memory_usage_percentage = self
+            .max_process_memory_usage_percentage
+            .map(|p| (self.system_physical_memory_bytes as f64 * p).round() as usize);
+        match (
+            self.max_process_memory_usage_bytes,
+            max_process_memory_usage_percentage,
+        ) {
+            (None, None) => None,
+            (Some(a), Some(b)) => Some(a.min(b)),
+            (Some(a), None) => Some(a),
+            (None, Some(b)) => Some(b),
+        }
+    }
+}
+
+impl Default for MemoryUsageBasedConnectionLimits {
+    fn default() -> Self {
+        Self::new()
+    }
+}

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

@@ -0,0 +1,41 @@
+// 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::fmt;
+
+/// A connection limit has been exceeded.
+#[derive(Debug, Clone, Copy)]
+pub struct MemoryUsageLimitExceeded {
+    pub process_physical_memory_bytes: usize,
+    pub max_allowed_bytes: usize,
+}
+
+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,
+        )
+    }
+}