Просмотр исходного кода

connector: add a stop signal to abort the Connector

Connector can be prematurely stopped even if a connection is in
progress.
draoi 2 лет назад
Родитель
Сommit
eca97916f6
2 измененных файлов с 40 добавлено и 6 удалено
  1. 3 0
      src/error.rs
  2. 37 6
      src/net/connector.rs

+ 3 - 0
src/error.rs

@@ -150,6 +150,9 @@ pub enum Error {
     #[error("Accept a new tls connection from the listener {0} failed")]
     AcceptTlsConnectionFailed(String),
 
+    #[error("Connector stopped")]
+    ConnectorStopped,
+
     #[error("Network operation failed")]
     NetworkOperationFailed,
 

+ 37 - 6
src/net/connector.rs

@@ -18,6 +18,10 @@
 
 use std::time::Duration;
 
+use futures::{
+    future::{select, Either},
+    pin_mut,
+};
 use log::warn;
 use url::Url;
 
@@ -28,7 +32,7 @@ use super::{
     settings::SettingsPtr,
     transport::Dialer,
 };
-use crate::{Error, Result};
+use crate::{system::CondVar, Error, Result};
 
 /// Create outbound socket connections
 pub struct Connector {
@@ -36,12 +40,14 @@ pub struct Connector {
     settings: SettingsPtr,
     /// Weak pointer to the session
     pub session: SessionWeakPtr,
+    /// Stop signal that aborts the connector if received.
+    stop_signal: CondVar,
 }
 
 impl Connector {
     /// Create a new connector with given network settings
     pub fn new(settings: SettingsPtr, session: SessionWeakPtr) -> Self {
-        Self { settings, session }
+        Self { settings, session, stop_signal: CondVar::new() }
     }
 
     /// Establish an outbound connection
@@ -72,10 +78,35 @@ impl Connector {
 
         let dialer = Dialer::new(endpoint.clone()).await?;
         let timeout = Duration::from_secs(self.settings.outbound_connect_timeout);
-        let ptstream = dialer.dial(Some(timeout)).await?;
 
-        let channel =
-            Channel::new(ptstream, Some(endpoint.clone()), url.clone(), self.session.clone()).await;
-        Ok((endpoint, channel))
+        let stop_fut = async {
+            self.stop_signal.wait().await;
+        };
+        let dial_fut = async { dialer.dial(Some(timeout)).await };
+
+        pin_mut!(stop_fut);
+        pin_mut!(dial_fut);
+
+        let result = {
+            match select(dial_fut, stop_fut).await {
+                Either::Left((Ok(ptstream), _)) => {
+                    let channel = Channel::new(
+                        ptstream,
+                        Some(endpoint.clone()),
+                        url.clone(),
+                        self.session.clone(),
+                    )
+                    .await;
+                    Ok((endpoint, channel))
+                }
+                Either::Left((Err(e), _)) => Err(e.into()),
+                Either::Right((_, _)) => return Err(Error::ConnectorStopped),
+            }
+        };
+        result
+    }
+
+    pub(crate) fn stop(&self) {
+        self.stop_signal.notify()
     }
 }