refactor: replace unbounded_channel with new_dialog_state_channel

Author: shenjinti

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "rsipstack"
-version = "0.2.83"
+version = "0.2.84"
 edition = "2021"
 description = "SIP Stack Rust library for building SIP applications"
 license = "MIT"

examples/client/main.rs

@@ -16,7 +16,6 @@ use rsipstack::{
 };
 use std::net::IpAddr;
 use std::{env, sync::Arc, time::Duration};
-use tokio::sync::mpsc::unbounded_channel;
 use tokio::time::timeout;
 use tokio::{select, time::sleep};
 use tokio_util::sync::CancellationToken;
@@ -212,7 +211,7 @@ async fn main() -> rsipstack::Result<()> {
     let incoming = endpoint.incoming_transactions()?;
     let dialog_layer = Arc::new(DialogLayer::new(endpoint.inner.clone()));
 
-    let (state_sender, state_receiver) = unbounded_channel();
+    let (state_sender, state_receiver) = dialog_layer.new_dialog_state_channel();
 
     let first_addr = endpoint
         .get_addrs()

src/bin/bench_ua.rs

@@ -23,7 +23,6 @@ use std::{
     },
     time::{Duration, Instant},
 };
-use tokio::sync::mpsc::unbounded_channel;
 use tokio::{select, time::sleep};
 use tokio_util::sync::CancellationToken;
 use tracing::{debug, info};
@@ -381,7 +380,7 @@ async fn main() -> Result<()> {
 
     let incoming = endpoint.incoming_transactions()?;
     let dialog_layer = Arc::new(DialogLayer::new(endpoint.inner.clone()));
-    let (state_sender, state_receiver) = unbounded_channel();
+    let (state_sender, state_receiver) = dialog_layer.new_dialog_state_channel();
     let stats = Stats::new();
 
     let mode_handler: BoxFuture<Result<()>> = match args.mode.as_str() {

src/dialog/dialog.rs

@@ -5,6 +5,7 @@ use super::{
     DialogId,
 };
 use crate::{
+    dialog::dialog_layer::DialogLayerInnerRef,
     rsip_ext::extract_uri_from_contact,
     transaction::{
         endpoint::EndpointInnerRef,
@@ -184,12 +185,75 @@ pub struct DialogInner {
     pub(super) initial_request: Request,
 }
 
-pub type DialogStateReceiver = UnboundedReceiver<DialogState>;
+pub struct DialogStateReceiver {
+    pub(super) dialog_layer_inner: DialogLayerInnerRef,
+    pub(super) receiver: UnboundedReceiver<DialogState>,
+    pub(super) dialog_id: Option<DialogId>,
+}
+
+impl DialogStateReceiver {
+    pub async fn recv(&mut self) -> Option<DialogState> {
+        let state = self.receiver.recv().await;
+        if let Some(ref s) = state {
+            if let Some(id) = &self.dialog_id {
+                if id != s.id() {
+                    match self.dialog_layer_inner.dialogs.write().as_mut() {
+                        Ok(dialogs) => {
+                            dialogs.remove(id);
+                        }
+                        Err(_) => {}
+                    }
+                }
+            }
+            self.dialog_id = Some(s.id().clone());
+        }
+        state
+    }
+}
+
+impl Drop for DialogStateReceiver {
+    fn drop(&mut self) {
+        let id = match self.dialog_id.take() {
+            Some(id) => id,
+            None => return,
+        };
+
+        match self.dialog_layer_inner.dialogs.write().as_mut() {
+            Ok(dialogs) => {
+                if let Some(dialog) = dialogs.remove(&id) {
+                    info!(%id, "dialog removed on state receiver drop");
+                    tokio::spawn(async move {
+                        if let Err(e) = dialog.hangup().await {
+                            warn!(%id, "error hanging up dialog on drop: {}", e);
+                        }
+                    });
+                }
+            }
+            Err(_) => {}
+        }
+    }
+}
+
 pub type DialogStateSender = UnboundedSender<DialogState>;
 
 pub(super) type DialogInnerRef = Arc<DialogInner>;
 
 impl DialogState {
+    pub fn id(&self) -> &DialogId {
+        match self {
+            DialogState::Calling(id)
+            | DialogState::Trying(id)
+            | DialogState::Early(id, _)
+            | DialogState::WaitAck(id, _)
+            | DialogState::Confirmed(id, _)
+            | DialogState::Updated(id, _)
+            | DialogState::Notify(id, _)
+            | DialogState::Info(id, _)
+            | DialogState::Options(id, _)
+            | DialogState::Terminated(id, _) => id,
+        }
+    }
+
     pub fn can_cancel(&self) -> bool {
         matches!(
             self,

src/dialog/dialog_layer.rs

@@ -1,7 +1,7 @@
 use super::authenticate::Credential;
 use super::dialog::DialogStateSender;
 use super::{dialog::Dialog, server_dialog::ServerInviteDialog, DialogId};
-use crate::dialog::dialog::DialogInner;
+use crate::dialog::dialog::{DialogInner, DialogStateReceiver};
 use crate::transaction::key::TransactionRole;
 use crate::transaction::make_tag;
 use crate::transaction::{endpoint::EndpointInnerRef, transaction::Transaction};
@@ -233,4 +233,14 @@ impl DialogLayer {
         let id = DialogId::try_from(req).ok()?;
         self.get_dialog(&id)
     }
+
+    pub fn new_dialog_state_channel(&self) -> (DialogStateSender, DialogStateReceiver) {
+        let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
+        let receiver = DialogStateReceiver {
+            dialog_layer_inner: self.inner.clone(),
+            receiver: rx,
+            dialog_id: None,
+        };
+        (tx, receiver)
+    }
 }

src/dialog/invitation.rs

@@ -5,7 +5,7 @@ use super::{
     dialog_layer::DialogLayer,
 };
 use crate::{
-    dialog::{dialog::Dialog, dialog_layer::DialogLayerInnerRef, DialogId},
+    dialog::{dialog::Dialog, DialogId},
     transaction::{
         key::{TransactionKey, TransactionRole},
         make_tag,
@@ -19,7 +19,7 @@ use rsip::{
     Request, Response,
 };
 use std::sync::Arc;
-use tracing::{debug, info, warn};
+use tracing::{debug, info};
 
 /// INVITE Request Options
 ///
@@ -134,64 +134,6 @@ pub struct InviteOption {
     pub headers: Option<Vec<rsip::Header>>,
 }
 
-pub struct DialogGuard {
-    pub dialog_layer_inner: DialogLayerInnerRef,
-    pub id: DialogId,
-}
-
-impl DialogGuard {
-    pub fn new(dialog_layer: &Arc<DialogLayer>, id: DialogId) -> Self {
-        Self {
-            dialog_layer_inner: dialog_layer.inner.clone(),
-            id,
-        }
-    }
-}
-
-impl Drop for DialogGuard {
-    fn drop(&mut self) {
-        let dlg = match self.dialog_layer_inner.dialogs.write() {
-            Ok(mut dialogs) => match dialogs.remove(&self.id) {
-                Some(dlg) => dlg,
-                None => return,
-            },
-            _ => return,
-        };
-        let _ = tokio::spawn(async move {
-            if let Err(e) = dlg.hangup().await {
-                info!(id=%dlg.id(), "failed to hangup dialog: {}", e);
-            }
-        });
-    }
-}
-
-pub(super) struct DialogGuardForUnconfirmed<'a> {
-    pub dialog_layer_inner: &'a DialogLayerInnerRef,
-    pub id: &'a DialogId,
-}
-
-impl<'a> Drop for DialogGuardForUnconfirmed<'a> {
-    fn drop(&mut self) {
-        // If the dialog is still unconfirmed, we should try to cancel it
-        match self.dialog_layer_inner.dialogs.write() {
-            Ok(mut dialogs) => match dialogs.remove(self.id) {
-                Some(dlg) => {
-                    info!(%self.id, "unconfirmed dialog dropped, cancelling it");
-                    let _ = tokio::spawn(async move {
-                        if let Err(e) = dlg.hangup().await {
-                            info!(id=%dlg.id(), "failed to hangup unconfirmed dialog: {}", e);
-                        }
-                    });
-                }
-                None => {}
-            },
-            Err(e) => {
-                warn!(%self.id, "failed to acquire write lock on dialogs: {}", e);
-            }
-        }
-    }
-}
-
 impl DialogLayer {
     /// Create an INVITE request from options
     ///
@@ -396,12 +338,6 @@ impl DialogLayer {
             .unwrap()
             .insert(id.clone(), Dialog::ClientInvite(dialog.clone()));
         info!(%id, "client invite dialog created");
-
-        let _guard = DialogGuardForUnconfirmed {
-            dialog_layer_inner: &self.inner,
-            id: &id,
-        };
-
         match dialog.process_invite(tx).await {
             Ok((new_dialog_id, resp)) => {
                 debug!(

src/dialog/server_dialog.rs

@@ -572,7 +572,6 @@ impl ServerInviteDialog {
             // discard old request
             return Ok(());
         }
-
         self.inner
             .remote_seq
             .compare_exchange(remote_seq, cseq, Ordering::Relaxed, Ordering::Relaxed)