rpc: Split r/w streams and lock them with Mutexes to allow better multiplexing

src/rpc/client.rs

@@ -86,20 +86,23 @@ impl RpcClient {
 
     /// Internal function that loops on a given stream and multiplexes the data
     async fn reqrep_loop(
-        mut stream: Box<dyn PtStream>,
+        stream: Box<dyn PtStream>,
         rep_send: channel::Sender<JsonResult>,
         req_recv: channel::Receiver<(JsonRequest, bool)>,
     ) -> Result<()> {
         debug!(target: "rpc::client::reqrep_loop()", "Starting reqrep loop");
+
+        let (mut reader, mut writer) = smol::io::split(stream);
+
         loop {
             let mut buf = Vec::with_capacity(INIT_BUF_SIZE);
 
             let (request, with_timeout) = req_recv.recv().await?;
 
             let request = JsonResult::Request(request);
-            write_to_stream(&mut stream, &request).await?;
+            write_to_stream(&mut writer, &request).await?;
 
-            let _ = read_from_stream(&mut stream, &mut buf, with_timeout).await?;
+            let _ = read_from_stream(&mut reader, &mut buf, with_timeout).await?;
             let val: JsonValue = String::from_utf8(buf)?.parse()?;
             let rep = JsonResult::try_from_value(&val)?;
             rep_send.send(rep).await?;

src/rpc/common.rs

@@ -18,7 +18,7 @@
 
 use std::time::Duration;
 
-use smol::io::{AsyncReadExt, AsyncWriteExt};
+use smol::io::{AsyncReadExt, AsyncWriteExt, ReadHalf, WriteHalf};
 
 use super::jsonrpc::*;
 use crate::{error::RpcError, net::transport::PtStream, system::io_timeout, Result};
@@ -29,7 +29,7 @@ pub(super) const READ_TIMEOUT: Duration = Duration::from_secs(30);
 
 /// Internal read function that reads from the active stream into a buffer.
 pub(super) async fn read_from_stream(
-    stream: &mut Box<dyn PtStream>,
+    reader: &mut ReadHalf<Box<dyn PtStream>>,
     buf: &mut Vec<u8>,
     with_timeout: bool,
 ) -> Result<usize> {
@@ -40,7 +40,7 @@ pub(super) async fn read_from_stream(
 
         // Lame we have to duplicate this code, but it is what it is.
         if with_timeout {
-            match io_timeout(READ_TIMEOUT, stream.read(&mut buf[total_read..])).await {
+            match io_timeout(READ_TIMEOUT, reader.read(&mut buf[total_read..])).await {
                 Ok(0) if total_read == 0 => {
                     return Err(
                         RpcError::ConnectionClosed("Connection closed cleanly".to_string()).into()
@@ -57,7 +57,7 @@ pub(super) async fn read_from_stream(
                 Err(e) => return Err(RpcError::IoError(e.kind()).into()),
             }
         } else {
-            match stream.read(&mut buf[total_read..]).await {
+            match reader.read(&mut buf[total_read..]).await {
                 Ok(0) if total_read == 0 => {
                     return Err(
                         RpcError::ConnectionClosed("Connection closed cleanly".to_string()).into()
@@ -83,7 +83,7 @@ pub(super) async fn read_from_stream(
 
 /// Internal write function that writes a JSON-RPC object to the active stream.
 pub(super) async fn write_to_stream(
-    stream: &mut Box<dyn PtStream>,
+    writer: &mut WriteHalf<Box<dyn PtStream>>,
     object: &JsonResult,
 ) -> Result<()> {
     let object_str = match object {
@@ -97,7 +97,7 @@ pub(super) async fn write_to_stream(
     // As we're a line-based protocol, we append the '\n' char at
     // the end of the JSON string.
     for i in [object_str.as_bytes(), &[b'\n']] {
-        if let Err(e) = stream.write_all(i).await {
+        if let Err(e) = writer.write_all(i).await {
             return Err(e.into())
         }
     }

src/rpc/server.rs

@@ -20,7 +20,10 @@ use std::{collections::HashSet, io::ErrorKind, sync::Arc};
 
 use async_trait::async_trait;
 use log::{debug, error, info};
-use smol::lock::MutexGuard;
+use smol::{
+    io::{ReadHalf, WriteHalf},
+    lock::{Mutex, MutexGuard},
+};
 use tinyjson::JsonValue;
 use url::Url;
 
@@ -73,7 +76,8 @@ pub trait RequestHandler: Sync + Send {
 /// Accept function that should run inside a loop for accepting incoming
 /// JSON-RPC requests and passing them to the [`RequestHandler`].
 pub async fn accept(
-    mut stream: Box<dyn PtStream>,
+    reader: Arc<Mutex<ReadHalf<Box<dyn PtStream>>>>,
+    writer: Arc<Mutex<WriteHalf<Box<dyn PtStream>>>>,
     addr: Url,
     rh: Arc<impl RequestHandler + 'static>,
     conn_limit: Option<usize>,
@@ -92,7 +96,11 @@ pub async fn accept(
 
     loop {
         let mut buf = Vec::with_capacity(INIT_BUF_SIZE);
-        let _ = read_from_stream(&mut stream, &mut buf, false).await?;
+
+        let mut reader_lock = reader.lock().await;
+        let _ = read_from_stream(&mut reader_lock, &mut buf, false).await?;
+        drop(reader_lock);
+
         let val: JsonValue = String::from_utf8(buf)?.trim().parse()?;
         let req = JsonRequest::try_from(&val)?;
 
@@ -111,17 +119,22 @@ pub async fn accept(
                     // Push notification
                     debug!(target: "rpc::server", "{} <-- {}", addr, notification.stringify()?);
                     let notification = JsonResult::Notification(notification);
-                    if let Err(e) = write_to_stream(&mut stream, &notification).await {
+
+                    let mut writer_lock = writer.lock().await;
+                    if let Err(e) = write_to_stream(&mut writer_lock, &notification).await {
                         subscription.unsubscribe().await;
                         return Err(e)
                     }
+                    drop(writer_lock);
                 }
             }
 
             JsonResult::SubscriberWithReply(subscriber, reply) => {
                 // Write the response
                 debug!(target: "rpc::server", "{} <-- {}", addr, reply.stringify()?);
-                write_to_stream(&mut stream, &reply.into()).await?;
+                let mut writer_lock = writer.lock().await;
+                write_to_stream(&mut writer_lock, &reply.into()).await?;
+                drop(writer_lock);
 
                 // Start the subscriber loop
                 let subscription = subscriber.sub.subscribe().await;
@@ -132,10 +145,14 @@ pub async fn accept(
                     // Push notification
                     debug!(target: "rpc::server", "{} <-- {}", addr, notification.stringify()?);
                     let notification = JsonResult::Notification(notification);
-                    if let Err(e) = write_to_stream(&mut stream, &notification).await {
+
+                    let mut writer_lock = writer.lock().await;
+                    if let Err(e) = write_to_stream(&mut writer_lock, &notification).await {
                         subscription.unsubscribe().await;
+                        drop(writer_lock);
                         return Err(e)
                     }
+                    drop(writer_lock);
                 }
             }
 
@@ -145,12 +162,16 @@ pub async fn accept(
 
             JsonResult::Response(ref v) => {
                 debug!(target: "rpc::server", "{} <-- {}", addr, v.stringify()?);
-                write_to_stream(&mut stream, &rep).await?;
+                let mut writer_lock = writer.lock().await;
+                write_to_stream(&mut writer_lock, &rep).await?;
+                drop(writer_lock);
             }
 
             JsonResult::Error(ref v) => {
                 debug!(target: "rpc::server", "{} <-- {}", addr, v.stringify()?);
-                write_to_stream(&mut stream, &rep).await?;
+                let mut writer_lock = writer.lock().await;
+                write_to_stream(&mut writer_lock, &rep).await?;
+                drop(writer_lock);
             }
         }
     }
@@ -169,10 +190,15 @@ async fn run_accept_loop(
             Ok((stream, url)) => {
                 let rh_ = rh.clone();
                 info!(target: "rpc::server", "[RPC] Server accepted conn from {}", url);
+
+                let (reader, writer) = smol::io::split(stream);
+                let reader = Arc::new(Mutex::new(reader));
+                let writer = Arc::new(Mutex::new(writer));
+
                 let task = StoppableTask::new();
                 let task_ = task.clone();
                 task.clone().start(
-                    accept(stream, url.clone(), rh.clone(), conn_limit),
+                    accept(reader, writer, url.clone(), rh.clone(), conn_limit),
                     |_| async move {
                         rh_.clone().unmark_connection(task_.clone()).await;
                     },