feat: implement client-side keepalive (#25)

Author: jjeffcaii

examples/cli.rs

@@ -1,9 +1,13 @@
+#[macro_use]
+extern crate log;
+
 use rsocket_rust::prelude::*;
+use rsocket_rust::Result;
 use rsocket_rust_transport_tcp::TcpClientTransport;
-use std::error::Error;
 
 #[tokio::main]
-async fn main() -> Result<(), Box<dyn Error + Send + Sync>> {
+async fn main() -> Result<()> {
+    env_logger::builder().format_timestamp_millis().init();
     let client = RSocketFactory::connect()
         .transport(TcpClientTransport::from("127.0.0.1:7878"))
         .acceptor(Box::new(|| {
@@ -15,8 +19,11 @@ async fn main() -> Result<(), Box<dyn Error + Send + Sync>> {
         .expect("Connect failed!");
 
     let req = Payload::builder().set_data_utf8("Ping!").build();
-    let res = client.request_response(req).await.expect("Requet failed!");
-    println!("request success: response={:?}", res);
+
+    match client.request_response(req).await {
+        Ok(res) => info!("{:?}", res),
+        Err(e) => error!("{}", e),
+    }
 
     Ok(())
 }

examples/echo.rs

@@ -4,14 +4,13 @@ extern crate log;
 use clap::{App, Arg, SubCommand};
 use rsocket_rust::prelude::*;
 use rsocket_rust::transport::Connection;
+use rsocket_rust::Result;
 use rsocket_rust_transport_tcp::{
     TcpClientTransport, TcpServerTransport, UnixClientTransport, UnixServerTransport,
 };
 use rsocket_rust_transport_websocket::{WebsocketClientTransport, WebsocketServerTransport};
 use std::fs;
 
-type Result<T> = rsocket_rust::Result<T>;
-
 enum RequestMode {
     FNF,
     REQUEST,

examples/proxy.rs

@@ -3,10 +3,9 @@ extern crate log;
 
 use futures::executor::block_on;
 use rsocket_rust::prelude::*;
+use rsocket_rust::Result;
 use rsocket_rust_transport_tcp::*;
 
-type Result<T> = rsocket_rust::Result<T>;
-
 #[tokio::main]
 async fn main() -> Result<()> {
     env_logger::builder().format_timestamp_millis().init();

examples/qps.rs

@@ -4,9 +4,9 @@ extern crate log;
 use clap::{App, Arg};
 use rsocket_rust::prelude::*;
 use rsocket_rust::transport::{Connection, Transport};
+use rsocket_rust::Result;
 use rsocket_rust_transport_tcp::{TcpClientTransport, UnixClientTransport};
 use rsocket_rust_transport_websocket::WebsocketClientTransport;
-use std::error::Error;
 use std::sync::{
     atomic::{AtomicU32, Ordering},
     Arc,
@@ -20,7 +20,7 @@ async fn connect<A, B>(
     count: u32,
     payload_size: usize,
     notify: Arc<Notify>,
-) -> Result<(), Box<dyn Error + Send + Sync>>
+) -> Result<()>
 where
     A: Send + Sync + Transport<Conn = B> + 'static,
     B: Send + Sync + Connection + 'static,
@@ -57,7 +57,7 @@ where
 }
 
 #[tokio::main]
-async fn main() -> Result<(), Box<dyn Error + Send + Sync>> {
+async fn main() -> Result<()> {
     env_logger::builder().format_timestamp_millis().init();
 
     let cli = App::new("echo")

rsocket/src/core/client.rs

@@ -1,4 +1,4 @@
-use crate::error::RSocketError;
+use crate::error::{RSocketError, ERR_CONN_CLOSED};
 use crate::frame::{self, Frame};
 use crate::payload::{Payload, SetupPayload, SetupPayloadBuilder};
 use crate::runtime;
@@ -7,7 +7,7 @@ use crate::transport::{
     self, Acceptor, Connection, DuplexSocket, Reader, Splitter, Transport, Writer,
 };
 use crate::Result;
-use futures::{future, FutureExt, SinkExt, StreamExt};
+use futures::{future, select, FutureExt, SinkExt, StreamExt};
 use std::error::Error;
 use std::net::SocketAddr;
 use std::pin::Pin;
@@ -115,6 +115,7 @@ where
         };
 
         let (snd_tx, mut snd_rx) = mpsc::channel::<Frame>(super::CHANNEL_SIZE);
+        let cloned_snd_tx = snd_tx.clone();
         let mut socket = DuplexSocket::new(1, snd_tx, splitter).await;
 
         let mut cloned_socket = socket.clone();
@@ -126,15 +127,40 @@ where
         let conn = tp.connect().await?;
         let (mut sink, mut stream) = conn.split();
 
+        let setup = self.setup.build();
+
+        // begin write loop
+        let tick_period = setup.keepalive_interval();
         runtime::spawn(async move {
-            while let Some(frame) = snd_rx.next().await {
-                if let Err(e) = (&mut sink).write(frame).await {
-                    error!("write frame failed: {}", e);
-                    break;
+            loop {
+                // send keepalive if timeout
+                match tokio::time::timeout(tick_period, snd_rx.next()).await {
+                    Ok(Some(frame)) => {
+                        if let frame::Body::Error(e) = frame.get_body_ref() {
+                            if e.get_code() == ERR_CONN_CLOSED {
+                                break;
+                            }
+                        }
+                        if let Err(e) = (&mut sink).write(frame).await {
+                            error!("write frame failed: {}", e);
+                            break;
+                        }
+                    }
+                    Ok(None) => break,
+                    Err(_) => {
+                        // keepalive
+                        let keepalive_frame =
+                            frame::Keepalive::builder(0, Frame::FLAG_RESPOND).build();
+                        if let Err(e) = (&mut sink).write(keepalive_frame).await {
+                            error!("write frame failed: {}", e);
+                            break;
+                        }
+                    }
                 }
             }
         });
 
+        // begin read loop
         let closer = self.closer.take();
         runtime::spawn(async move {
             while let Some(next) = stream.read().await {
@@ -151,12 +177,22 @@ where
                     }
                 }
             }
+
+            // workaround: send a notify frame that the connection has been closed.
+            let close_frame = frame::Error::builder(0, 0)
+                .set_code(ERR_CONN_CLOSED)
+                .build();
+            if let Err(_) = cloned_snd_tx.send(close_frame).await {
+                debug!("send close notify frame failed!");
+            }
+
+            // invoke on_close handler
             if let Some(mut invoke) = closer {
                 invoke();
             }
         });
 
-        socket.setup(self.setup.build()).await;
+        socket.setup(setup).await;
         Ok(Client::from(socket))
     }
 }

rsocket/src/frame/mod.rs

@@ -187,6 +187,10 @@ impl Frame {
         self.body
     }
 
+    pub fn get_body_ref(&self) -> &Body {
+        &self.body
+    }
+
     pub fn get_flag(&self) -> u16 {
         self.flag
     }