From db76cc3844ccf642d3842353f8a7e59d6316fc21 Mon Sep 17 00:00:00 2001
From: Thomas Eizinger <thomas@eizinger.io>
Date: Tue, 1 Apr 2025 11:08:07 +1100
Subject: [PATCH] fix(relay): reduce memory usage of eBPF program to < 100MB
 (#8587)

At present, the eBPF program would try to pre-allocate around 800MB of
memory for all entries in the maps. This would allow for 1 million
channel mappings. We don't need that many to begin with. Reducing the
max number of channels down to 65536 reduces our memory usage to less
than 100MB.

Related: #7518

---------

Signed-off-by: Thomas Eizinger <thomas@eizinger.io>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
---
 rust/relay/ebpf-turn-router/src/main.rs | 36 ++++++++++++++++++++++---
 1 file changed, 32 insertions(+), 4 deletions(-)

diff --git a/rust/relay/ebpf-turn-router/src/main.rs b/rust/relay/ebpf-turn-router/src/main.rs
index f6d5e9d6e..5e4aa5ec6 100644
--- a/rust/relay/ebpf-turn-router/src/main.rs
+++ b/rust/relay/ebpf-turn-router/src/main.rs
@@ -36,21 +36,23 @@ mod slice_mut_at;
 mod stats;
 mod udp;
 
+const NUM_ENTRIES: u32 = 0x10000;
+
 /// Channel mappings from an IPv4 socket + channel number to an IPv4 socket + port.
 ///
 /// TODO: Update flags to `BPF_F_NO_PREALLOC` to guarantee atomicity? Needs research.
 #[map]
 static CHAN_TO_UDP_44: HashMap<ClientAndChannelV4, PortAndPeerV4> =
-    HashMap::with_max_entries(0x100000, 0);
+    HashMap::with_max_entries(NUM_ENTRIES, 0);
 #[map]
 static UDP_TO_CHAN_44: HashMap<PortAndPeerV4, ClientAndChannelV4> =
-    HashMap::with_max_entries(0x100000, 0);
+    HashMap::with_max_entries(NUM_ENTRIES, 0);
 #[map]
 static CHAN_TO_UDP_66: HashMap<ClientAndChannelV6, PortAndPeerV6> =
-    HashMap::with_max_entries(0x100000, 0);
+    HashMap::with_max_entries(NUM_ENTRIES, 0);
 #[map]
 static UDP_TO_CHAN_66: HashMap<PortAndPeerV6, ClientAndChannelV6> =
-    HashMap::with_max_entries(0x100000, 0);
+    HashMap::with_max_entries(NUM_ENTRIES, 0);
 
 #[xdp]
 pub fn handle_turn(ctx: XdpContext) -> u32 {
@@ -296,3 +298,29 @@ fn on_panic(_: &core::panic::PanicInfo) -> ! {
 fn main() {
     panic!("This program is meant to be compiled as an eBPF program.");
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// Memory overhead of an eBPF map.
+    ///
+    /// Determined empirically.
+    const HASH_MAP_OVERHEAD: f32 = 1.5;
+
+    #[test]
+    fn hashmaps_are_less_than_100_mb() {
+        let ipv4_datatypes =
+            core::mem::size_of::<PortAndPeerV4>() + core::mem::size_of::<ClientAndChannelV4>();
+        let ipv6_datatypes =
+            core::mem::size_of::<PortAndPeerV6>() + core::mem::size_of::<ClientAndChannelV6>();
+
+        let ipv4_map_size = ipv4_datatypes as f32 * NUM_ENTRIES as f32 * HASH_MAP_OVERHEAD;
+        let ipv6_map_size = ipv6_datatypes as f32 * NUM_ENTRIES as f32 * HASH_MAP_OVERHEAD;
+
+        let total_map_size = (ipv4_map_size + ipv6_map_size) * 2_f32;
+        let total_map_size_mb = total_map_size / 1024_f32 / 1024_f32;
+
+        assert!(total_map_size_mb < 100_f32);
+    }
+}