feat(socket): fix some stuff, add blocking connect to ReqSocket

Author: mempirate

msg-socket/src/req/driver.rs

@@ -15,7 +15,7 @@ use tokio::{
     time::Interval,
 };
 use tokio_util::codec::Framed;
-use tracing::{debug, error, trace};
+use tracing::{debug, error, trace, warn};
 
 use super::{Command, ReqError, ReqOptions};
 use crate::{ConnectionState, ExponentialBackoff, req::SocketState};
@@ -327,7 +327,7 @@ where
                 }
                 Poll::Ready(Some(Err(err))) => {
                     if let reqrep::Error::Io(e) = err {
-                        error!(err = ?e, "Socket wire error");
+                        error!(err = ?e, "wire error, resetting connection state");
                     }
 
                     // set the connection to inactive, so that it will be re-tried
@@ -336,9 +336,12 @@ where
                     continue;
                 }
                 Poll::Ready(None) => {
-                    debug!("Connection to {:?} closed, shutting down driver", this.addr);
+                    warn!(peer = ?this.addr, "connection closed, resetting connection state");
 
-                    return Poll::Ready(());
+                    // set the connection to inactive, so that it will be re-tried
+                    this.reset_connection();
+
+                    continue;
                 }
                 Poll::Pending => {}
             }

msg-socket/src/req/mod.rs

@@ -35,6 +35,8 @@ pub enum ReqError {
     Timeout,
     #[error("Could not connect to any valid endpoints")]
     NoValidEndpoints,
+    #[error("Failed to connect to the target endpoint: {0:?}")]
+    Connect(Box<dyn std::error::Error + Send + Sync>),
 }
 
 /// Commands that can be sent to the request socket driver.

msg-socket/src/req/socket.rs

@@ -115,7 +115,8 @@ where
         // Initialize communication channels
         let (to_driver, from_socket) = mpsc::channel(DEFAULT_BUFFER_SIZE);
 
-        let transport = self.transport.take().expect("Transport has been moved already");
+        // TODO: Don't panic, return error
+        let mut transport = self.transport.take().expect("Transport has been moved already");
 
         // We initialize the connection as inactive, and let it be activated
         // by the backend task as soon as the driver is spawned.
@@ -124,6 +125,13 @@ where
             backoff: ExponentialBackoff::new(Duration::from_millis(20), 16),
         };
 
+        if self.options.blocking_connect {
+            transport
+                .connect(endpoint.clone())
+                .await
+                .map_err(|e| ReqError::Connect(Box::new(e)))?;
+        }
+
         let timeout_check_interval = tokio::time::interval(self.options.timeout / 10);
 
         let flush_interval = self.options.flush_interval.map(tokio::time::interval);

msg-socket/tests/it/reqrep.rs

@@ -1,3 +1,5 @@
+use std::time::Duration;
+
 use bytes::Bytes;
 use msg_socket::{RepSocket, ReqSocket};
 use msg_transport::{
@@ -144,3 +146,70 @@ async fn reqrep_mutual_tls_works() {
     let response = req.request(hello.clone()).await.unwrap();
     assert_eq!(hello, response, "expected {:?}, got {:?}", hello, response);
 }
+
+#[tokio::test]
+async fn reqrep_late_bind_works() {
+    let _ = tracing_subscriber::fmt::try_init();
+
+    let mut rep = RepSocket::new(Tcp::default());
+    let mut req = ReqSocket::new(Tcp::default());
+
+    let local_addr = "localhost:64521";
+    req.connect(local_addr).await.unwrap();
+
+    let hello = Bytes::from_static(b"hello");
+
+    let reply = tokio::spawn(async move { req.request(hello.clone()).await.unwrap() });
+
+    tokio::time::sleep(Duration::from_millis(1000)).await;
+    rep.bind(local_addr).await.unwrap();
+
+    let msg = rep.next().await.unwrap();
+    let payload = msg.msg().clone();
+    msg.respond(payload).unwrap();
+
+    let response = reply.await.unwrap();
+    let hello = Bytes::from_static(b"hello");
+    assert_eq!(hello, response, "expected {:?}, got {:?}", hello, response);
+}
+
+#[tokio::test]
+async fn reqrep_drop_server() {
+    let _ = tracing_subscriber::fmt::try_init();
+
+    let mut rep = RepSocket::new(Tcp::default());
+    let mut req = ReqSocket::new(Tcp::default());
+
+    rep.bind("0.0.0.0:0").await.unwrap();
+
+    let addr = rep.local_addr().unwrap().clone();
+    req.connect(addr).await.unwrap();
+
+    tokio::spawn(async move {
+        let request = rep.next().await.unwrap();
+        let msg = request.msg().clone();
+        request.respond(msg).unwrap();
+
+        drop(rep);
+    });
+
+    let hello = Bytes::from_static(b"hello");
+    let response = req.request(hello.clone()).await.unwrap();
+    assert_eq!(hello, response, "expected {:?}, got {:?}", hello, response);
+
+    match req.request(hello.clone()).await {
+        Ok(response) => assert_eq!(hello, response, "expected {:?}, got {:?}", hello, response),
+        Err(e) => tracing::warn!("Error: {:?}", e),
+    }
+
+    tokio::time::sleep(Duration::from_secs(60)).await;
+
+    tokio::spawn(async move {
+        req.request(hello.clone()).await.unwrap();
+    });
+
+    let mut rep = RepSocket::new(Tcp::default());
+    rep.bind(addr).await.unwrap();
+
+    tokio::time::sleep(Duration::from_millis(10000)).await;
+}