RSDK-828 - fix file handle leak (#7)

Co-authored-by: Nicolas Menard <[email protected]>

Author: stuqdog

src/ffi/dial_ffi.rs

@@ -11,15 +11,16 @@ use tokio::sync::oneshot;
 use tracing::Level;
 
 use crate::rpc::dial::{
-    CredentialsExt, DialBuilder, DialOptions, WithCredentials, WithoutCredentials,
+    CredentialsExt, DialBuilder, DialOptions, ViamChannel, WithCredentials, WithoutCredentials,
 };
 use libc::c_char;
 
 use crate::proxy;
 use hyper::Server;
 use std::ffi::{CStr, CString};
-use tower::{make::Shared, ServiceBuilder};
+use tower::{make::Shared, util::Either, ServiceBuilder};
 use tower_http::{
+    auth::AddAuthorization,
     trace::{DefaultMakeSpan, DefaultOnRequest, DefaultOnResponse, TraceLayer},
     LatencyUnit,
 };
@@ -32,6 +33,7 @@ use crate::proxy::grpc_proxy::GRPCProxy;
 pub struct DialFfi {
     runtime: Option<Runtime>,
     sigs: Option<Vec<oneshot::Sender<()>>>,
+    channels: Vec<Either<AddAuthorization<ViamChannel>, ViamChannel>>,
 }
 
 impl Drop for DialFfi {
@@ -48,6 +50,7 @@ impl DialFfi {
         Self {
             runtime: Some(Runtime::new().unwrap()),
             sigs: None,
+            channels: vec![],
         }
     }
     fn push_signal(&mut self, sig: oneshot::Sender<()>) {
@@ -172,18 +175,19 @@ pub unsafe extern "C" fn dial(
     } else {
         disable_webrtc = uri_str.contains(".local") || uri_str.contains("localhost");
     }
-    let server = match runtime.block_on(async move {
-        let dial = match payload {
-            Some(p) => tower::util::Either::A(
+    let (server, channel) = match runtime.block_on(async move {
+        let channel = match payload {
+            Some(p) => Either::A(
                 dial_with_cred(uri_str, p.to_str()?, allow_insec, disable_webrtc)?
                     .connect()
                     .await?,
             ),
             None => {
                 let c = dial_without_cred(uri_str, allow_insec, disable_webrtc)?;
-                tower::util::Either::B(c.connect().await?)
+                Either::B(c.connect().await?)
             }
         };
+        let dial = channel.clone();
         let g = GRPCProxy::new(dial, uri);
         let service = ServiceBuilder::new()
             .layer(
@@ -200,14 +204,15 @@ pub unsafe extern "C" fn dial(
         let server = Server::builder(conn)
             .http2_only(true)
             .serve(Shared::new(service));
-        Ok::<_, Box<dyn std::error::Error>>(server)
+        Ok::<_, Box<dyn std::error::Error>>((server, channel))
     }) {
         Ok(s) => s,
         Err(e) => {
             println!("Error building GRPC proxy reason : {e:?}");
             return ptr::null_mut();
         }
     };
+    ctx.channels.push(channel);
     let server = server.with_graceful_shutdown(async {
         rx.await.ok();
     });
@@ -248,13 +253,25 @@ pub extern "C" fn free_rust_runtime(rt_ptr: Option<Box<DialFfi>>) -> i32 {
             return -1;
         }
     };
-    match ctx.sigs.take() {
-        Some(sigs) => {
-            for sig in sigs {
-                let _ = sig.send(());
-            }
+    if let Some(sigs) = ctx.sigs.take() {
+        for sig in sigs {
+            let _ = sig.send(());
+        }
+    }
+
+    for channel in &ctx.channels {
+        let channel = match channel {
+            Either::A(chan) => chan.get_ref(),
+            Either::B(chan) => chan,
+        };
+        match channel {
+            ViamChannel::Direct(_) => (),
+            ViamChannel::WebRTC(chan) => ctx
+                .runtime
+                .as_ref()
+                .map(|rt| rt.block_on(async move { chan.close().await }))
+                .unwrap_or_default(),
         }
-        None => {}
     }
     log::debug!("Freeing rust runtime");
     0

src/ffi/spatialmath/vector3.rs

@@ -13,42 +13,42 @@ pub unsafe extern "C" fn free_vector_memory(ptr: *mut Vector3) {
     if ptr.is_null() {
         return;
     }
-    Box::from_raw(ptr);
+    let _ = Box::from_raw(ptr);
 }
 
 #[no_mangle]
 pub unsafe extern "C" fn vector_get_components(vec_ptr: *const Vector3) -> *const c_double {
     null_pointer_check!(vec_ptr);
     let vec = *vec_ptr;
-    let components: [c_double;3] = [vec.x, vec.y, vec.z];
+    let components: [c_double; 3] = [vec.x, vec.y, vec.z];
     Box::into_raw(Box::new(components)) as *const _
 }
 
 #[no_mangle]
 pub unsafe extern "C" fn vector_set_x(vec_ptr: *mut Vector3, x_val: f64) {
     null_pointer_check!(vec_ptr);
-    let vec = &mut*vec_ptr;
+    let vec = &mut *vec_ptr;
     vec.x = x_val
 }
 
 #[no_mangle]
 pub unsafe extern "C" fn vector_set_y(vec_ptr: *mut Vector3, y_val: f64) {
     null_pointer_check!(vec_ptr);
-    let vec = &mut*vec_ptr;
+    let vec = &mut *vec_ptr;
     vec.y = y_val
 }
 
 #[no_mangle]
 pub unsafe extern "C" fn vector_set_z(vec_ptr: *mut Vector3, z_val: f64) {
     null_pointer_check!(vec_ptr);
-    let vec = &mut*vec_ptr;
+    let vec = &mut *vec_ptr;
     vec.z = z_val
 }
 
 #[no_mangle]
 pub unsafe extern "C" fn normalize_vector(vec_ptr: *mut Vector3) {
     null_pointer_check!(vec_ptr);
-    let vec = &mut*vec_ptr;
+    let vec = &mut *vec_ptr;
     vec.normalize();
 }
 
@@ -63,7 +63,7 @@ pub unsafe extern "C" fn vector_get_normalized(vec_ptr: *const Vector3) -> *mut
 #[no_mangle]
 pub unsafe extern "C" fn scale_vector(vec_ptr: *mut Vector3, factor: f64) {
     null_pointer_check!(vec_ptr);
-    let vec = &mut*vec_ptr;
+    let vec = &mut *vec_ptr;
     vec.scale(factor);
 }
 
@@ -77,7 +77,8 @@ pub unsafe extern "C" fn vector_get_scaled(vec_ptr: *const Vector3, factor: f64)
 
 #[no_mangle]
 pub unsafe extern "C" fn vector_add(
-    vec_ptr_1: *const Vector3, vec_ptr_2: *const Vector3
+    vec_ptr_1: *const Vector3,
+    vec_ptr_2: *const Vector3,
 ) -> *mut Vector3 {
     null_pointer_check!(vec_ptr_1);
     null_pointer_check!(vec_ptr_2);
@@ -88,7 +89,8 @@ pub unsafe extern "C" fn vector_add(
 
 #[no_mangle]
 pub unsafe extern "C" fn vector_subtract(
-    vec_ptr_1: *const Vector3, vec_ptr_2: *const Vector3
+    vec_ptr_1: *const Vector3,
+    vec_ptr_2: *const Vector3,
 ) -> *mut Vector3 {
     null_pointer_check!(vec_ptr_1);
     null_pointer_check!(vec_ptr_2);
@@ -98,7 +100,10 @@ pub unsafe extern "C" fn vector_subtract(
 }
 
 #[no_mangle]
-pub unsafe extern "C" fn vector_dot_product(vec_ptr_1: *const Vector3, vec_ptr_2: *const Vector3) -> f64 {
+pub unsafe extern "C" fn vector_dot_product(
+    vec_ptr_1: *const Vector3,
+    vec_ptr_2: *const Vector3,
+) -> f64 {
     null_pointer_check!(vec_ptr_1, f64::NAN);
     null_pointer_check!(vec_ptr_2, f64::NAN);
     let vec1 = &*vec_ptr_1;
@@ -108,12 +113,13 @@ pub unsafe extern "C" fn vector_dot_product(vec_ptr_1: *const Vector3, vec_ptr_2
 
 #[no_mangle]
 pub unsafe extern "C" fn vector_cross_product(
-    vec_ptr_1: *mut Vector3, vec_ptr_2: *mut Vector3
+    vec_ptr_1: *mut Vector3,
+    vec_ptr_2: *mut Vector3,
 ) -> *mut Vector3 {
     null_pointer_check!(vec_ptr_1);
     null_pointer_check!(vec_ptr_2);
-    let vec1 = &mut*vec_ptr_1;
-    let vec2 = &mut*vec_ptr_2;
+    let vec1 = &mut *vec_ptr_1;
+    let vec2 = &mut *vec_ptr_2;
     let vec = vec1.cross(vec2);
     vec.to_raw_pointer()
 }

src/rpc/base_channel.rs

@@ -1,7 +1,10 @@
 use anyhow::Result;
-use std::sync::{
-    atomic::{AtomicBool, AtomicPtr, Ordering},
-    Arc,
+use std::{
+    fmt::Debug,
+    sync::{
+        atomic::{AtomicBool, AtomicPtr, Ordering},
+        Arc,
+    },
 };
 use webrtc::{data_channel::RTCDataChannel, peer_connection::RTCPeerConnection};
 
@@ -14,16 +17,28 @@ pub struct WebRTCBaseChannel {
     closed: AtomicBool,
 }
 
+impl Debug for WebRTCBaseChannel {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("WebRTCBaseChannel")
+            .field("Peer connection id", &self.peer_connection.get_stats_id())
+            .field("Data channel id", &self.data_channel.id())
+            .finish()
+    }
+}
+
 impl WebRTCBaseChannel {
     pub(crate) async fn new(
         peer_connection: Arc<RTCPeerConnection>,
         data_channel: Arc<RTCDataChannel>,
     ) -> Arc<Self> {
         let dc = data_channel.clone();
-        let pc = peer_connection.clone();
+        let pc = Arc::downgrade(&peer_connection);
         peer_connection
             .on_ice_connection_state_change(Box::new(move |conn_state| {
-                let pc = pc.clone();
+                let pc = match pc.upgrade(){
+                    Some(pc) => pc,
+                    None => return Box::pin(async  {}),
+                };
                 Box::pin(async move {
                     let sctp = pc.sctp();
                     let transport = sctp.transport();
@@ -44,9 +59,12 @@ impl WebRTCBaseChannel {
             closed: AtomicBool::new(false),
         });
 
-        let c = channel.clone();
+        let c = Arc::downgrade(&channel);
         dc.on_error(Box::new(move |err: webrtc::Error| {
-            let c = c.clone();
+            let c = match c.upgrade() {
+                Some(c) => c,
+                None => return Box::pin(async {}),
+            };
             Box::pin(async move {
                 if let Err(e) = c.close_with_reason(Some(anyhow::Error::from(err))).await {
                     log::error!("error closing channel: {e}")

src/rpc/client_channel.rs

@@ -7,9 +7,12 @@ use anyhow::Result;
 use chashmap::CHashMap;
 use hyper::Body;
 use prost::Message;
-use std::sync::{
-    atomic::{AtomicBool, AtomicPtr, AtomicU64, Ordering},
-    Arc,
+use std::{
+    fmt::Debug,
+    sync::{
+        atomic::{AtomicBool, AtomicPtr, AtomicU64, Ordering},
+        Arc,
+    },
 };
 use webrtc::{
     data_channel::{data_channel_message::DataChannelMessage, RTCDataChannel},
@@ -27,7 +30,35 @@ pub struct WebRTCClientChannel {
     pub(crate) receiver_bodies: CHashMap<u64, hyper::Body>,
 }
 
+impl Debug for WebRTCClientChannel {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("WebRTCClientChannel")
+            .field("stream_id_counter", &self.stream_id_counter)
+            .field("base channel", &self.base_channel)
+            .finish()
+    }
+}
+
+impl Drop for WebRTCClientChannel {
+    fn drop(&mut self) {
+        let bc = self.base_channel.clone();
+        if !bc.is_closed() {
+            let _ = tokio::spawn(async move {
+                if let Err(e) = bc.close().await {
+                    log::error!("Error closing base channel: {e}");
+                }
+            });
+        };
+        log::debug!("Dropping client channel {:?}", &self);
+    }
+}
+
 impl WebRTCClientChannel {
+    pub async fn close(&self) {
+        self.base_channel.close().await.unwrap();
+        self.base_channel.data_channel.close().await.unwrap();
+        self.base_channel.peer_connection.close().await.unwrap();
+    }
     pub(crate) async fn new(
         peer_connection: Arc<RTCPeerConnection>,
         data_channel: Arc<RTCDataChannel>,
@@ -42,17 +73,25 @@ impl WebRTCClientChannel {
 
         let channel = Arc::new(channel);
         let ret_channel = channel.clone();
+        let channel = Arc::downgrade(&channel);
 
         data_channel
             .on_message(Box::new(move |msg: DataChannelMessage| {
                 let channel = channel.clone();
                 Box::pin(async move {
+                    let channel = match channel.upgrade() {
+                        Some(channel) => channel,
+                        None => {
+                            return;
+                        }
+                    };
                     if let Err(e) = channel.on_channel_message(msg).await {
                         log::error!("error deserializing message: {e}");
                     }
                 })
             }))
             .await;
+        log::debug!("Client channel created");
         ret_channel
     }