Bugfix/fix turn unit test memory leak (#626)

mutexd · web-flow · commit 03237d0a3689 · 2024-11-03T13:03:22.000-08:00
* Fix cyclic dependency in ChannelBind/Permission

- Both structs have a timer-callback like mechanism to remove itself from the table which held by Allocation. In this case, cyclic dependency is created when the struct holds Arc&lt;&gt; pointing to the table owning it. We replace the Arc&lt;&gt; in the callback with Weak&lt;&gt;

* Fix cyclic dependency in Allocation

- Allocation hold Arc&lt;&gt; to the table of Allocations owned by AllocationManager and thus creates the cyclic dependency. We replace it with Weak&lt;&gt; when pointing to parent
diff --git a/turn/src/allocation/allocation_manager.rs b/turn/src/allocation/allocation_manager.rs
@@ -20,7 +20,7 @@ pub struct ManagerConfig {
 
 /// `Manager` is used to hold active allocations.
 pub struct Manager {
-    allocations: AllocationMap,
+    allocations: Arc<Mutex<AllocationMap>>,
     reservations: Arc<Mutex<HashMap<String, u16>>>,
     relay_addr_generator: Box<dyn RelayAddressGenerator + Send + Sync>,
     alloc_close_notify: Option<mpsc::Sender<AllocationInfo>>,
@@ -107,9 +107,9 @@ impl Manager {
             relay_addr,
             five_tuple,
             username,
+            Arc::downgrade(&self.allocations),
             self.alloc_close_notify.clone(),
         );
-        a.allocations = Some(Arc::clone(&self.allocations));
 
         log::debug!("listening on relay addr: {:?}", a.relay_addr);
         a.start(lifetime).await;
diff --git a/turn/src/allocation/allocation_test.rs b/turn/src/allocation/allocation_test.rs
@@ -12,12 +12,14 @@ async fn test_has_permission() -> Result<()> {
     let turn_socket = Arc::new(UdpSocket::bind("0.0.0.0:0").await?);
     let relay_socket = Arc::clone(&turn_socket);
     let relay_addr = relay_socket.local_addr()?;
+    let allocations = Arc::new(Mutex::new(AllocationMap::new()));
     let a = Allocation::new(
         turn_socket,
         relay_socket,
         relay_addr,
         FiveTuple::default(),
         TextAttribute::new(ATTR_USERNAME, "user".into()),
+        Arc::downgrade(&allocations),
         None,
     );
 
@@ -50,12 +52,14 @@ async fn test_add_permission() -> Result<()> {
     let turn_socket = Arc::new(UdpSocket::bind("0.0.0.0:0").await?);
     let relay_socket = Arc::clone(&turn_socket);
     let relay_addr = relay_socket.local_addr()?;
+    let allocations = Arc::new(Mutex::new(AllocationMap::new()));
     let a = Allocation::new(
         turn_socket,
         relay_socket,
         relay_addr,
         FiveTuple::default(),
         TextAttribute::new(ATTR_USERNAME, "user".into()),
+        Arc::downgrade(&allocations),
         None,
     );
 
@@ -74,12 +78,14 @@ async fn test_remove_permission() -> Result<()> {
     let turn_socket = Arc::new(UdpSocket::bind("0.0.0.0:0").await?);
     let relay_socket = Arc::clone(&turn_socket);
     let relay_addr = relay_socket.local_addr()?;
+    let allocations = Arc::new(Mutex::new(AllocationMap::new()));
     let a = Allocation::new(
         turn_socket,
         relay_socket,
         relay_addr,
         FiveTuple::default(),
         TextAttribute::new(ATTR_USERNAME, "user".into()),
+        Arc::downgrade(&allocations),
         None,
     );
 
@@ -107,12 +113,14 @@ async fn test_add_channel_bind() -> Result<()> {
     let turn_socket = Arc::new(UdpSocket::bind("0.0.0.0:0").await?);
     let relay_socket = Arc::clone(&turn_socket);
     let relay_addr = relay_socket.local_addr()?;
+    let allocations = Arc::new(Mutex::new(AllocationMap::new()));
     let a = Allocation::new(
         turn_socket,
         relay_socket,
         relay_addr,
         FiveTuple::default(),
         TextAttribute::new(ATTR_USERNAME, "user".into()),
+        Arc::downgrade(&allocations),
         None,
     );
 
@@ -141,12 +149,14 @@ async fn test_get_channel_by_number() -> Result<()> {
     let turn_socket = Arc::new(UdpSocket::bind("0.0.0.0:0").await?);
     let relay_socket = Arc::clone(&turn_socket);
     let relay_addr = relay_socket.local_addr()?;
+    let allocations = Arc::new(Mutex::new(AllocationMap::new()));
     let a = Allocation::new(
         turn_socket,
         relay_socket,
         relay_addr,
         FiveTuple::default(),
         TextAttribute::new(ATTR_USERNAME, "user".into()),
+        Arc::downgrade(&allocations),
         None,
     );
 
@@ -177,12 +187,14 @@ async fn test_get_channel_by_addr() -> Result<()> {
     let turn_socket = Arc::new(UdpSocket::bind("0.0.0.0:0").await?);
     let relay_socket = Arc::clone(&turn_socket);
     let relay_addr = relay_socket.local_addr()?;
+    let allocations = Arc::new(Mutex::new(AllocationMap::new()));
     let a = Allocation::new(
         turn_socket,
         relay_socket,
         relay_addr,
         FiveTuple::default(),
         TextAttribute::new(ATTR_USERNAME, "user".into()),
+        Arc::downgrade(&allocations),
         None,
     );
 
@@ -209,12 +221,14 @@ async fn test_remove_channel_bind() -> Result<()> {
     let turn_socket = Arc::new(UdpSocket::bind("0.0.0.0:0").await?);
     let relay_socket = Arc::clone(&turn_socket);
     let relay_addr = relay_socket.local_addr()?;
+    let allocations = Arc::new(Mutex::new(AllocationMap::new()));
     let a = Allocation::new(
         turn_socket,
         relay_socket,
         relay_addr,
         FiveTuple::default(),
         TextAttribute::new(ATTR_USERNAME, "user".into()),
+        Arc::downgrade(&allocations),
         None,
     );
 
@@ -246,12 +260,14 @@ async fn test_allocation_refresh() -> Result<()> {
     let turn_socket = Arc::new(UdpSocket::bind("0.0.0.0:0").await?);
     let relay_socket = Arc::clone(&turn_socket);
     let relay_addr = relay_socket.local_addr()?;
+    let allocations = Arc::new(Mutex::new(AllocationMap::new()));
     let a = Allocation::new(
         turn_socket,
         relay_socket,
         relay_addr,
         FiveTuple::default(),
         TextAttribute::new(ATTR_USERNAME, "user".into()),
+        Arc::downgrade(&allocations),
         None,
     );
 
@@ -268,12 +284,14 @@ async fn test_allocation_close() -> Result<()> {
     let turn_socket = Arc::new(UdpSocket::bind("0.0.0.0:0").await?);
     let relay_socket = Arc::clone(&turn_socket);
     let relay_addr = relay_socket.local_addr()?;
+    let allocations = Arc::new(Mutex::new(AllocationMap::new()));
     let a = Allocation::new(
         turn_socket,
         relay_socket,
         relay_addr,
         FiveTuple::default(),
         TextAttribute::new(ATTR_USERNAME, "user".into()),
+        Arc::downgrade(&allocations),
         None,
     );
 
diff --git a/turn/src/allocation/channel_bind.rs b/turn/src/allocation/channel_bind.rs
@@ -2,7 +2,7 @@
 mod channel_bind_test;
 
 use std::sync::atomic::Ordering;
-use std::sync::Arc;
+use std::sync::{Arc, Weak};
 
 use portable_atomic::AtomicBool;
 use tokio::sync::Mutex;
@@ -18,7 +18,7 @@ use crate::proto::channum::*;
 pub struct ChannelBind {
     pub(crate) peer: SocketAddr,
     pub(crate) number: ChannelNumber,
-    pub(crate) channel_bindings: Option<Arc<Mutex<HashMap<ChannelNumber, ChannelBind>>>>,
+    pub(crate) channel_bindings: Option<Weak<Mutex<HashMap<ChannelNumber, ChannelBind>>>>,
     reset_tx: Option<mpsc::Sender<Duration>>,
     timer_expired: Arc<AtomicBool>,
 }
@@ -51,7 +51,7 @@ impl ChannelBind {
             while !done {
                 tokio::select! {
                     _ = &mut timer => {
-                        if let Some(cbs) = &channel_bindings{
+                        if let Some(cbs) = &channel_bindings.clone().and_then(|x| x.upgrade()) {
                             let mut cb = cbs.lock().await;
                             if cb.remove(&number).is_none() {
                                 log::error!("Failed to remove ChannelBind for {}", number);
diff --git a/turn/src/allocation/channel_bind/channel_bind_test.rs b/turn/src/allocation/channel_bind/channel_bind_test.rs
@@ -12,12 +12,14 @@ async fn create_channel_bind(lifetime: Duration) -> Result<Allocation> {
     let turn_socket = Arc::new(UdpSocket::bind("0.0.0.0:0").await?);
     let relay_socket = Arc::clone(&turn_socket);
     let relay_addr = relay_socket.local_addr()?;
+    let allocations = Arc::new(Mutex::new(AllocationMap::new()));
     let a = Allocation::new(
         turn_socket,
         relay_socket,
         relay_addr,
         FiveTuple::default(),
         TextAttribute::new(ATTR_USERNAME, "user".into()),
+        Arc::downgrade(&allocations),
         None,
     );
 
diff --git a/turn/src/allocation/mod.rs b/turn/src/allocation/mod.rs
@@ -10,7 +10,7 @@ use std::collections::HashMap;
 use std::marker::{Send, Sync};
 use std::net::SocketAddr;
 use std::sync::atomic::Ordering;
-use std::sync::Arc;
+use std::sync::{Arc, Weak};
 
 use channel_bind::*;
 use five_tuple::*;
@@ -34,7 +34,7 @@ use crate::proto::*;
 
 const RTP_MTU: usize = 1500;
 
-pub type AllocationMap = Arc<Mutex<HashMap<FiveTuple, Arc<Allocation>>>>;
+pub type AllocationMap = HashMap<FiveTuple, Arc<Allocation>>;
 
 /// Information about an [`Allocation`].
 #[derive(Debug, Clone)]
@@ -77,7 +77,7 @@ pub struct Allocation {
     username: Username,
     permissions: Arc<Mutex<HashMap<String, Permission>>>,
     channel_bindings: Arc<Mutex<HashMap<ChannelNumber, ChannelBind>>>,
-    pub(crate) allocations: Option<AllocationMap>,
+    allocations: Weak<Mutex<AllocationMap>>,
     reset_tx: SyncMutex<Option<mpsc::Sender<Duration>>>,
     timer_expired: Arc<AtomicBool>,
     closed: AtomicBool, // Option<mpsc::Receiver<()>>,
@@ -98,6 +98,7 @@ impl Allocation {
         relay_addr: SocketAddr,
         five_tuple: FiveTuple,
         username: Username,
+        allocation_map: Weak<Mutex<AllocationMap>>,
         alloc_close_notify: Option<mpsc::Sender<AllocationInfo>>,
     ) -> Self {
         Allocation {
@@ -109,7 +110,7 @@ impl Allocation {
             username,
             permissions: Arc::new(Mutex::new(HashMap::new())),
             channel_bindings: Arc::new(Mutex::new(HashMap::new())),
-            allocations: None,
+            allocations: allocation_map,
             reset_tx: SyncMutex::new(None),
             timer_expired: Arc::new(AtomicBool::new(false)),
             closed: AtomicBool::new(false),
@@ -137,7 +138,7 @@ impl Allocation {
             }
         }
 
-        p.permissions = Some(Arc::clone(&self.permissions));
+        p.permissions = Some(Arc::downgrade(&self.permissions));
         p.start(PERMISSION_TIMEOUT).await;
 
         {
@@ -184,7 +185,7 @@ impl Allocation {
         let peer = c.peer;
 
         // Add or refresh this channel.
-        c.channel_bindings = Some(Arc::clone(&self.channel_bindings));
+        c.channel_bindings = Some(Arc::downgrade(&self.channel_bindings));
         c.start(lifetime).await;
 
         {
@@ -279,7 +280,7 @@ impl Allocation {
             while !done {
                 tokio::select! {
                     _ = &mut timer => {
-                        if let Some(allocs) = &allocations{
+                        if let Some(allocs) = &allocations.upgrade(){
                             let mut allocs = allocs.lock().await;
                             if let Some(a) = allocs.remove(&five_tuple) {
                                 let _ = a.close().await;
@@ -355,7 +356,7 @@ impl Allocation {
                         match result {
                             Ok((n, src_addr)) => (n, src_addr),
                             Err(_) => {
-                                if let Some(allocs) = &allocations {
+                                if let Some(allocs) = &allocations.upgrade() {
                                     let mut allocs = allocs.lock().await;
                                     allocs.remove(&five_tuple);
                                 }
diff --git a/turn/src/allocation/permission.rs b/turn/src/allocation/permission.rs
@@ -1,5 +1,5 @@
 use std::sync::atomic::Ordering;
-use std::sync::Arc;
+use std::sync::{Arc, Weak};
 
 use portable_atomic::AtomicBool;
 use tokio::sync::Mutex;
@@ -15,7 +15,7 @@ pub(crate) const PERMISSION_TIMEOUT: Duration = Duration::from_secs(5 * 60);
 /// https://tools.ietf.org/html/rfc5766#section-2.3
 pub struct Permission {
     pub(crate) addr: SocketAddr,
-    pub(crate) permissions: Option<Arc<Mutex<HashMap<String, Permission>>>>,
+    pub(crate) permissions: Option<Weak<Mutex<HashMap<String, Permission>>>>,
     reset_tx: Option<mpsc::Sender<Duration>>,
     timer_expired: Arc<AtomicBool>,
 }
@@ -47,7 +47,7 @@ impl Permission {
             while !done {
                 tokio::select! {
                     _ = &mut timer => {
-                        if let Some(perms) = &permissions{
+                        if let Some(perms) = &permissions.clone().and_then(|x| x.upgrade()) {
                             let mut p = perms.lock().await;
                             p.remove(&addr2ipfingerprint(&addr));
                         }
